Auto-import all attention modules and export their registered classes.

Sub-modules

neuralops.attentions.adaptive_span_attention

Adaptive Span Attention — Final Corrected Implementation. Based on: Sukhbaatar et al., 2019 (arXiv:1905.07799) …

neuralops.attentions.adaptive_sparse_attention

Adaptive Sparse Attention Module (Production-Hardened). Implements dynamic top-k and threshold-based adaptive sparse attention with LoRA adapters, …

neuralops.attentions.adaptive_svd_attention

Adaptive SVD Attention — Production-Hardened (v2) …

neuralops.attentions.adasplash_attention

AdaSplash Attention: Adaptive Splash Attention with α-Entmax Sparse Activations …

neuralops.attentions.aegis_hybrid

Aegis-Hybrid Attention — Production-Hardened (v2) …

neuralops.attentions.aegis_spectron_fp8

Aegis-Spectron-FP8 Hybrid Attention (Final Corrected) …

neuralops.attentions.aegis_v2

Aegis-Attention V2 (Corrected) …

neuralops.attentions.aegis_v5

AEGIS-Attention V5 (Corrected): Adaptive Enhanced Global-Local Integrated Sparse Attention …

neuralops.attentions.aegis_v6

AEGIS-Attention V6 (Final-Final Corrected).

neuralops.attentions.asvd_attention

ASVD (Activation-Aware Singular Value Decomposition) Attention. Paper: https://arxiv.org/abs/2307.07382 …

neuralops.attentions.bam_umoe

BAM-UMoE: Unified Mixture of Experts Attention Block (Activation-Aware v2) …

neuralops.attentions.block_transformer_hiba

Block Transformer with Hierarchical Intra-Block and Inter-Block Attention (HIBA) …

neuralops.attentions.chela_slay

CHELA-SLAY: Spherical Linearized Attention with Yat-Kernel — Production-Hardened (v3) …

neuralops.attentions.cluster_attention

ClusterFormer: Token Clustering Attention Module …

neuralops.attentions.compressed_attention

Compressed Multi-Head Attention via Truncated Singular Value Decomposition (SVD) …

neuralops.attentions.compressed_mha_attention

Compressed Multi-Head Attention for DistilBERT and Multi-Architecture via Joint Truncated SVD …

neuralops.attentions.cosformer

Cosformer: Rethinking Softmax in Attention with Linear Complexity …

neuralops.attentions.decay_linear_attention

Decay Linear Attention (DecayLinearAttention) …

neuralops.attentions.eat

EAT: Efficient Attention Transformer with Progressive Token Pruning. Production-Hardened (v2) …

neuralops.attentions.expert_choice_moe

Expert Choice Mixture of Experts (MoE) Attention Block …

neuralops.attentions.factorised_matrix_state

Factorised Matrix State Attention Module …

neuralops.attentions.flash_attention3

FlashAttention-3: Fast and Memory-Efficient Exact Attention with Asynchronous Hardware Overlap …

neuralops.attentions.flash_attention3_fp8

FlashAttention-3 FP8: Low-Precision 8-bit Float Transformer Attention …

neuralops.attentions.full_attention

Standard Multi-Head Scaled Dot-Product Attention (Vaswani et al., 2017) …

neuralops.attentions.fused_svd_kernel

Fused SVD-Linear Layer with SRAM-Resident RMSNorm & Dual-Projection …

neuralops.attentions.fusion_attn

FUSIONATTN: Multi-Paradigm Hybrid Sparse Attention Module (Production-Hardened v2) …

neuralops.attentions.galerkin_attention

Galerkin Attention: Petrov-Galerkin Projection Attention (Corrected v3) …

neuralops.attentions.gated_delta_net

Gated DeltaNet Attention Module — Production-Hardened (v2).

neuralops.attentions.gated_deltanet_block

Gated DeltaNet Block with Optional Multi-Head Attention Bypass …

neuralops.attentions.gated_linear_attention

Gated Linear Attention (GLA) — Production-Hardened (v3) …

neuralops.attentions.goldfinch_ahn

GoldFinch-AHN: Adaptive Hybrid Network — Sliding Window + Gated Recurrent Memory. Production-Hardened (v2).

neuralops.attentions.gqa_attention

Grouped-Query Attention (GQA) — v3, rank-aware …

neuralops.attentions.granite_hybrid

Granite Hybrid Attention-SSM Layer with MoE Feed-Forward …

neuralops.attentions.griffin_ssd

Griffin SSD: Hybrid Local Sliding-Window Attention and Global State Space Duality (SSD) Kernel …

neuralops.attentions.halo_regla

HALO-ReGLA: Hardware-Aware Linear Optics & Rectified Gated Linear Attention – Production-Hardened …

neuralops.attentions.hamming_attention

Hamming Attention Distillation (HAD) Module (Corrected) …

neuralops.attentions.hgrn2_gated_deltanet

HGRN-2 Gated DeltaNet Hybrid Attention Module — Production-Hardened (v2) …

neuralops.attentions.hgrn2_layer

Hierarchically Gated Recurrent Network V2 (HGRN-2) Layer …

neuralops.attentions.hgrn2_outer

HGRN-2 Outer-Product Recurrent Linear Attention — Production-Hardened v2 …

neuralops.attentions.hla_attention

Hadamard Linear Attention (HLA) …

neuralops.attentions.hybrid_attention

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

neuralops.attentions.hymba_transx

HymbaTransX: Hybrid Differential Transformer & State Space Model Block (Corrected) …

neuralops.attentions.infini_attention

Infini-Attention with Compressive Memory and Sliding Window — Fixed (v3) …

neuralops.attentions.jamba

Jamba-2 Hybrid Transformer: SVD-Mamba + MHA + MoE Attention Block .

neuralops.attentions.jamba_info

Information-Bottlenecked Jamba Architecture with Dual Operational Modes (v2 — Production-Hardened) …

neuralops.attentions.kda_attention

Kernelized Decay Attention (KDA) / Recurrent Linear Attention (Corrected) …

neuralops.attentions.leka_v1

LEKAv1: Linear-time Efficient Kernel Attention (Corrected with Positive Features) …

neuralops.attentions.leka_v2

LEKAv2: Linear-time Efficient Kernel Attention with Landmark Distillation …

neuralops.attentions.leka_v3

LEKAv3: Linear-time Efficient Kernel Attention with Value-Scored Landmarks (Corrected O(N)) …

neuralops.attentions.linformer_attention

Linformer Attention – Production-hardened with full multi-architecture support.

neuralops.attentions.lola_attention

LoLA: Low-Rank Linear Attention with Sparse Caching …

neuralops.attentions.lora_attention

LoRA Parameter-Efficient Fine-Tuning Attention Module (v2, Production-Hardened) …

neuralops.attentions.lsh_attention

Locality-Sensitive Hashing (LSH) Attention – Reformer style (Corrected) …

neuralops.attentions.mamba_block

Faithful Mamba Block with SVD-based Attention Weight Distillation …

neuralops.attentions.mamba_svd_flash

MambaSVD-Flash v3: Hybrid Bidirectional SSM-FlashAttention Module …

neuralops.attentions.minimal_nova

neuralops.attentions.minimal_svd

neuralops.attentions.mixattention_moska

MixAttention-MoSKA: Frequency/Variance-Aware Sparse Attention (FASA). Production-Hardened (v2).

neuralops.attentions.mla_attention

Multi-Head Latent Attention (MLA) – DeepSeek-V2 (Production-Hardened v2) …

neuralops.attentions.moba_attention

MoBA: Mixture of Block Attention — Production-Hardened (v2) …

neuralops.attentions.moe_deltakv_paged

MoE-DeltaKV-Paged Hybrid Attention — Production-Hardened v2 …

neuralops.attentions.naive_svd

Naive SVD baseline: plain truncated SVD without calibration, Fisher weighting, or whitening. Used as an uncalibrated low-rank baseline in comparative …

neuralops.attentions.nasc_attention

NASC: NEXUS-AEGIS Spectral Cascade Attention. Hardened for multi-architecture production deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5 …

neuralops.attentions.nexus

NEXUS: Neural Efficient X-scoring Unified Sparse Attention …

neuralops.attentions.nha_attention

Native Hybrid Attention (NHA) – Production-Hardened (v2) …

neuralops.attentions.nsa_attention

Native Sparse Attention (NSA) Module (Production-Hardened) …

neuralops.attentions.nystrom_attention

Nyströmformer Attention – Production-Hardened …

neuralops.attentions.optimized_hybrid_attention

Optimized Hybrid Attention (OHA) — Production-Hardened (v2) …

neuralops.attentions.orion

ORION: Optimized Rectified Input-aware Orthogonal Network (Corrected) …

neuralops.attentions.padre_attention

PADRe: Polynomial Attention with Dropout Regularization. Replaces softmax with a stable polynomial approximation, enabling faster inference and …

neuralops.attentions.performer_attention

Performer Attention – Fast Attention via Positive Orthogonal Random Features (FAVOR+) …

neuralops.attentions.phantom

Phantom: Parameter-Hashed AttentioN with Token-wise Orthogonal Mapping – Production-Hardened …

neuralops.attentions.primal_attention

Primal Attention: Self-Attention through Asymmetric Kernel SVD in Primal Representation …

neuralops.attentions.progressive_token_pruning

Progressive Token Pruning (PTP) – Production-Hardened (v2) …

neuralops.attentions.projection_factorized_attn

Projection Factorized SVD Linear Attention — Production-Hardened (v2) …

neuralops.attentions.qcqa_attention

Quality and Capacity-Aware Grouped Query Attention (QCQA) — v2 …

neuralops.attentions.rala_attention

Rank-Augmented Linear Attention (RALA) – Corrected …

neuralops.attentions.reap_moe

REAP-MoE: Robust Expert Allocation and Pruning Mixture of Experts Attention Module …

neuralops.attentions.regla_attention

Refined Gated Linear Attention (ReGLA) – Production-Hardened (v2).

neuralops.attentions.replace

Model Attention Replacement and Dynamic Graph Surgery Utilities — v2.

neuralops.attentions.res_svd

Spectral Residual Low-Rank SVD Attention (Res-SVD Attention) – Production-Hardened …

neuralops.attentions.routing_attention

Routing Attention: Content-Based Clustering for Sparse Attention – Production-Hardened …

neuralops.attentions.rpc_attention

Random Projection Compressed (RPC) Attention — Production-Hardened (v2) …

neuralops.attentions.salt_attention

SALT Attention: Singular Value Adaptation with Low-Rank Tuning …

neuralops.attentions.samba_m3

SAMBA-M3 v2: Hybrid Sliding-Window Attention and Spectral Mamba Block (Hardened).

neuralops.attentions.samba_moe

Samba-MoE: Hybrid Selective State Space (SAMBA-M3) + REAP-MoE Block. Combines local sliding-window attention, spectral Mamba SSM, and MoE routing. …

neuralops.attentions.selective_ssm

Selective State Space Model (SelectiveSSM) …

neuralops.attentions.selective_ssm

Selective State Space Model (SelectiveSSM) …

neuralops.attentions.soa_attn

Shared Orthogonal Attention (SOA) with Per-Projection Latent Bottlenecks. Hardened for production use across diverse transformer architectures: - …

neuralops.attentions.span_expanded_attention

Span-Expanded Attention with HyLoRA – Production-Hardened …

neuralops.attentions.spectra_v2

MHA-SDPA: A thin Multi-Head Attention wrapper around PyTorch's scaled_dot_product_attention …

neuralops.attentions.spectraformer_attention

Spectraformer: Random Feature Linear Attention with Flexible Kernel Approximation. Replaces softmax attention with random orthogonal feature …

neuralops.attentions.spectral_mamba

Spectral-Mamba Attention Module: Hybrid Attention-State Space Module with Spectral Channel Modulation – Production-Hardened …

neuralops.attentions.spectron

SpectronAttention: Importance-Aware Low-Rank + Sparse Attention (Activation-Aware v2) …

neuralops.attentions.squeezing_heads_distill

Squeezing Heads Attention Distillation Layer (v2, Real Head Compression) …

neuralops.attentions.svd_attention

SVD-Compressed Sequence Attention …

neuralops.attentions.svd_compressed

SVD-Compressed Attention with Importance-Aware Lagrangian Rank Allocation …

neuralops.attentions.svd_linear_attention

SVD Factorized Linear Attention & Kernelized Linear Attention with Exponential Decay Gating — Production-Hardened (v2) …

neuralops.attentions.svd_linear_decay

SVD Linear Decay Gated Attention Alias. Maps to framework.attentions.svd_linear_attention.SVDLinearAttention.

neuralops.attentions.svd_llm_v2_attention

SVD-LLM V2 Attention: Truncation-Aware SVD for LLM Compression. Direct low-rank factorized parameterization maintaining explicit U, Sigma, and V^H …

neuralops.attentions.svd_mamba

SVD-Mamba: SVD-factorized Selective State Space Model. Hardened for production use across diverse transformer architectures: - DistilBERT (Flat MHA) - …

neuralops.attentions.svd_projection_attention

SVD Projection Attention. Each projection (Q, K, V, Out) is factored into low-rank SVD components U and V, with trainable LoRA adapters for recovery …

neuralops.attentions.template_attention

Template Attention Mechanism (YourCustomAttention) – Production-Hardened. Supports: DistilBERT (Flat MHA), RoBERTa/BERT (Nested Encoder), GPT-2 …

neuralops.attentions.tlinformer_attention

TLinFormer: Exact linear-time attention via low-rank sequence projections. Reference: LinFormer (Wang et al., 2020) …

neuralops.attentions.token_compressed_attention

Token-Compressed Attention — Production-Hardened (v2) …

neuralops.attentions.tpa_attention

TPAttention: Tensor Product Attention with Factorised Q-K Projections …

neuralops.attentions.triton_cosformer

Triton-Accelerated Cosformer Linear Attention Kernel …

neuralops.attentions.triton_svd_linear_attention

Triton-Accelerated SVD Linear Attention & Chunked Causal Linear Attention …

neuralops.attentions.unified_rope

Universal Rotary Position Embedding Attention (UnifiedRoPE). Hardened for production use across diverse transformer architectures: - DistilBERT (Flat …

neuralops.attentions.upi_stuffed_mamba

UPI-Stuffed Mamba: Uncertainty-Penalized Interpolated Mamba SSM Block. Hardened for production use across diverse transformer architectures: - …

neuralops.attentions.vig_attention

ViG Attention: Vision GNN-inspired 2D Grid Linear Attention with Gated Fusion. Hardened for production use across diverse transformer architectures: - …

neuralops.attentions.vspec

VSpec++ Attention: Variational Spectral Token-Adaptive Attention. Hardened for production use across diverse transformer architectures: - DistilBERT …

neuralops.attentions.wgqa_attention

Weighted Grouped-Query Attention (WGQA) …

neuralops.attentions.xattention

XAttention: Block Sparse Attention with Coarse-to-Fine Routing …

neuralops.attentions.zeros_attention

Zero-Mean Linear Attention (ZeroS Attention) …

Classes

class ASVDAttention (config: Any | None = None,
rank: int | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "asvd", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "asvd_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class ASVDAttention(BaseAttention):
    r"""Hardened ASVD Attention with Factored Low-Rank Matrix Decomposition, GQA, RoPE, and LoRA.

    Computational Complexity:
        - Factorized Projections: $\mathcal{O}(B \cdot N \cdot r \cdot (d + d_{\text{all}}))$
        - Attention Scoring: $\mathcal{O}(B \cdot H \cdot N^2 \cdot d_k)$

    Tensor Dimensions:
        - hidden_states: $[B, N, D]$ where $B$ is batch size, $N$ is sequence length, and $D$ is hidden dimension.
        - Output: Tuple of $([B, N, D], \text{None})$.
    """

    is_custom_attention: bool = True
    triton_kernel: str = "low_rank"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__()
        self.config = config
        model = config
        self.hidden_size: int = kwargs.get(
            "hidden_size",
            kwargs.get(
                "d_model",
                getattr(
                    model,
                    "hidden_size",
                    getattr(getattr(model, "config", None), "hidden_size", 64),
                ),
            ),
        )
        self.num_heads: int = kwargs.get(
            "num_heads",
            kwargs.get(
                "num_attention_heads",
                getattr(
                    model,
                    "num_attention_heads",
                    getattr(getattr(model, "config", None), "num_attention_heads", 4),
                ),
            ),
        )
        self.num_kv_heads: int = kwargs.get(
            "num_kv_heads",
            kwargs.get(
                "num_key_value_heads",
                getattr(
                    model,
                    "num_key_value_heads",
                    getattr(
                        getattr(model, "config", None),
                        "num_key_value_heads",
                        self.num_heads,
                    ),
                ),
            ),
        )
        self.head_dim: int = kwargs.get(
            "head_dim",
            getattr(
                model,
                "head_dim",
                getattr(getattr(model, "config", None), "head_dim", None),
            ),
        ) or (self.hidden_size // max(1, self.num_heads))
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))

        # Rank configuration
        self.rank: int = getattr(
            config,
            "asvd_rank",
            (rank if rank is not None else getattr(config, "rank", kwargs.pop("rank", 32))),
        )
        self.lora_rank: int = getattr(config, "lora_rank", kwargs.get("lora_rank", 32))
        self.freeze_base: bool = getattr(
            config, "freeze_base", kwargs.get("freeze_base", self.lora_rank > 0)
        )
        self.energy_threshold: float = getattr(
            config, "asvd_energy_threshold", kwargs.get("energy_threshold", 0.95)
        )
        self.error_threshold: float | None = getattr(
            config, "asvd_error_threshold", kwargs.get("error_threshold", 0.40)
        )
        self.dynamic_rank: bool = getattr(
            config, "asvd_dynamic_rank", kwargs.get("dynamic_rank", True)
        )

        # Flags & metadata
        self.supports_gqa: bool = True
        self.supports_rope: bool = True
        self.is_nested_encoder: bool = False
        self.rotary_emb: Any | None = None

        # Decoder detection
        model_type = getattr(config, "model_type", "").lower()
        is_dec = getattr(config, "is_decoder", False)
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(model_type.startswith(p) for p in decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)

        # Projections
        self.query = ASVDLinear(
            self.hidden_size,
            self.all_head_size,
            self.rank,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=self.freeze_base,
        )
        self.key = ASVDLinear(
            self.hidden_size,
            self.kv_head_size,
            self.rank,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=self.freeze_base,
        )
        self.value = ASVDLinear(
            self.hidden_size,
            self.kv_head_size,
            self.rank,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=self.freeze_base,
        )
        self.out_proj = ASVDLinear(
            self.all_head_size,
            self.hidden_size,
            self.rank,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=self.freeze_base,
        )

        dropout_p = getattr(
            config,
            "dropout",
            getattr(
                config,
                "hidden_dropout_prob",
                getattr(config, "attention_probs_dropout_prob", 0.1),
            ),
        )
        self.dropout = nn.Dropout(dropout_p)
        self.triton_kwargs = {}

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

    # --------------------------------------------------------------------------
    # Compatibility properties for parameter inspection & framework integration
    # --------------------------------------------------------------------------
    @property
    def q_proj(self) -> ASVDLinear:
        return self.query

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

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

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

    @property
    def q_lin(self) -> ASVDLinear:
        return self.query

    @property
    def k_lin(self) -> ASVDLinear:
        return self.key

    @property
    def v_lin(self) -> ASVDLinear:
        return self.value

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

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

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

    def _decompose_and_set(
        self,
        target_linear: ASVDLinear,
        src_weight: torch.Tensor,
        src_bias: torch.Tensor | None,
        proj_name: str = "",
    ) -> float:
        """Decomposes src_weight into low-rank SVD factors and verifies relative Frobenius error."""
        with torch.no_grad():
            W = src_weight.detach().float()
            # Ensure correct orientation: [out_features, in_features]
            if W.ndim != 2:
                raise ValueError(
                    f"Expected 2D weight matrix for projection {proj_name}, got shape {W.shape}"
                )
            if (
                W.shape[0] != target_linear.out_features
                and W.shape[1] == target_linear.out_features
            ):
                W = W.t().contiguous()

            # Perform SVD strictly in FP32
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)

            # Energy-based dynamic rank selection
            total_energy = torch.sum(S**2)
            cum_energy = torch.cumsum(S**2, dim=0) / (total_energy + 1e-12)

            target_energy = max(
                self.energy_threshold,
                (1.0 - (self.error_threshold**2) if self.error_threshold is not None else 0.95),
            )
            indices = (cum_energy >= target_energy).nonzero()
            if len(indices) > 0:
                r_energy = indices[0].item() + 1
            else:
                r_energy = min(U.size(1), Vt.size(0))

            # Cap rank so factorized layer strictly compresses or matches original parameter count
            # Dense size: in_features * out_features. Factored size: r * (in_features + out_features) + r
            max_param_r = max(
                1,
                int(
                    (target_linear.in_features * target_linear.out_features)
                    / (target_linear.in_features + target_linear.out_features + 1)
                ),
            )

            if self.dynamic_rank:
                r = max(16, min(r_energy, U.size(1), Vt.size(0), max_param_r))
            else:
                r = min(self.rank, U.size(1), Vt.size(0), max_param_r)

            w_rec = (U[:, :r] * S[:r].unsqueeze(0)) @ Vt[:r, :]
            diff = torch.norm(W - w_rec, p="fro")
            denom = torch.norm(W, p="fro") + 1e-9
            rel_err = float((diff / denom).item())

            # If error exceeds threshold under dynamic rank, expand rank within safe compression bound
            if self.error_threshold is not None and rel_err > self.error_threshold:
                if self.dynamic_rank and r < min(U.size(1), Vt.size(0), max_param_r):
                    indices_strict = (
                        cum_energy >= (1.0 - (self.error_threshold**2) + 1e-5)
                    ).nonzero()
                    if len(indices_strict) > 0:
                        r = min(indices_strict[0].item() + 1, max_param_r)
                        w_rec = (U[:, :r] * S[:r].unsqueeze(0)) @ Vt[:r, :]
                        rel_err = float((torch.norm(W - w_rec) / denom).item())

                if rel_err > self.error_threshold:
                    logger.warning(
                        f"Relative Frobenius reconstruction error ({rel_err:.2%}) for projection '{proj_name}' "
                        f"exceeds {self.error_threshold:.0%} threshold. Proceeding with best-effort rank {r}."
                    )

            target_linear.set_svd_factors(U[:, :r], Vt[:r, :], S[:r], bias=src_bias)
            return rel_err

    def copy_weights(self, src: Any) -> None:
        """Extracts and factorizes pretrained weights across DistilBERT, RoBERTa, GPT-2, Qwen/LLaMA."""
        extracted_weights: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. Robust Nested Encoder Detection (prior to extraction)
        # Robustly detects both:
        # a) Outer RobertaAttention / BertAttention with inner 'self' (query/key/value) and 'output.dense'
        # b) Inner RobertaSelfAttention / BertSelfAttention
        # Prevents double output projection signal attenuation
        cls_name = src.__class__.__name__
        has_out_lin = hasattr(src, "out_lin")

        is_nested_encoder = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            and hasattr(src, "output")
            and hasattr(src.output, "dense")
        ):
            # Outer RobertaAttention containing inner self-attention and output dense
            is_nested_encoder = True
        elif (
            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:
            # Inner SelfAttention module
            is_nested_encoder = True

        self.is_nested_encoder = bool(is_nested_encoder)

        # 2. GPT-2: Fused Conv1D c_attn + Conv1D c_proj
        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()
                c_b = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                c_b = c_b.data.clone() if c_b is not None else None
                c_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in c_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    c_w = c_w.t().contiguous()
                split_size = c_w.size(0) // 3
                extracted_weights["q"] = (
                    c_w[:split_size, :],
                    c_b[:split_size] if c_b is not None else None,
                )
                extracted_weights["k"] = (
                    c_w[split_size : 2 * split_size, :],
                    c_b[split_size : 2 * split_size] if c_b is not None else None,
                )
                extracted_weights["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()
                    o_b = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                    o_b = o_b.data.clone() if o_b is not None else None
                    p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                    if (
                        ("Conv1D" in p_cls)
                        or hasattr(base_o, "nf")
                        or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                        or hasattr(c_proj_mod, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        o_w = o_w.t().contiguous()
                    extracted_weights["o"] = (o_w, o_b)

        # 3. RoBERTa / BERT Nested Encoder: self.query, self.key, self.value
        s = getattr(src, "self", src)
        name_map = {"q": "query", "k": "key", "v": "value"}
        for proj, attr in name_map.items():
            if (
                extracted_weights[proj][0] is None
                and hasattr(s, attr)
                and hasattr(getattr(s, attr), "weight")
            ):
                mod = getattr(s, attr)
                b = getattr(mod, "bias", None)
                extracted_weights[proj] = (
                    mod.weight.data,
                    b.data if b is not None else None,
                )

        if (
            not self.is_nested_encoder
            and extracted_weights["o"][0] is None
            and hasattr(src, "output")
            and hasattr(src.output, "dense")
        ):
            mod = src.output.dense
            b = getattr(mod, "bias", None)
            extracted_weights["o"] = (
                mod.weight.data,
                b.data if b is not None else None,
            )

        # 4. Qwen / LLaMA: q_proj, k_proj, v_proj, o_proj
        for proj in ["q", "k", "v", "o"]:
            if extracted_weights[proj][0] is None:
                for candidate in [f"{proj}_proj", f"{proj}_lin", proj]:
                    if hasattr(src, candidate) and hasattr(getattr(src, candidate), "weight"):
                        mod = getattr(src, candidate)
                        b = getattr(mod, "bias", None)
                        extracted_weights[proj] = (
                            mod.weight.data,
                            b.data if b is not None else None,
                        )
                        break

        # 5. Fallback to UniversalAdapter
        for proj in ["q", "k", "v", "o"]:
            if extracted_weights[proj][0] is None:
                try:
                    w, b = UniversalAdapter.get_weight_bias(src, proj)
                    if w is not None:
                        extracted_weights[proj] = (w, b)
                except Exception as e:
                    print(f"copy_weights failed: {e}")
                    logger.warning(f"copy_weights failed: {e}")

        # 6. Check for missing essential projections and Frobenius norm verification
        missing = [p for p in ["q", "k", "v"] if extracted_weights[p][0] is None]
        if not self.is_nested_encoder and extracted_weights["o"][0] is None:
            missing.append("o")
        if missing:
            raise RuntimeError(
                f"ASVDAttention.copy_weights: could not extract mandatory projections {missing} "
                f"from {src.__class__.__name__}. Refusing to install randomly-initialised weights."
            )

        for p in ["q", "k", "v"]:
            w_check = extracted_weights[p][0]
            if w_check is not None:
                norm = float(torch.norm(w_check.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        f"ASVDAttention extracted weight for '{p}' has near-zero Frobenius norm ({norm:.2e})"
                    )

        if extracted_weights["o"][0] is not None:
            norm_o = float(torch.norm(extracted_weights["o"][0].float()).item())
            if norm_o < 1e-4:
                logger.warning(
                    f"ASVDAttention extracted weight for 'o' has near-zero Frobenius norm ({norm_o:.2e})"
                )

        # 7. SVD Factorization & Reconstruction Verification
        for proj, target in [
            ("q", self.query),
            ("k", self.key),
            ("v", self.value),
            ("o", self.out_proj),
        ]:
            if proj == "o" and self.is_nested_encoder:
                self.out_proj.set_identity()
                continue
            w, b = extracted_weights[proj]
            if w is not None:
                self._decompose_and_set(target, w, b, proj_name=proj)

        # 8. Inherit Decoder & Rotary Metadata
        if hasattr(src, "rotary_emb"):
            self.rotary_emb = src.rotary_emb
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)
        elif hasattr(src, "c_attn") or cls_name in (
            "GPT2Attention",
            "Qwen2Attention",
            "Qwen2_5_VLAttention",
            "QwenAttention",
        ):
            self.is_decoder = True

    def forward(
        self,
        hidden_states: torch.Tensor | None = None,
        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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Forward pass for factorized ASVD attention."""
        if hidden_states is None:
            hidden_states = kwargs.get("query", kwargs.get("x"))
        if hidden_states is None and len(args) > 0 and isinstance(args[0], torch.Tensor):
            hidden_states = args[0]
        if attention_mask is None:
            attention_mask = kwargs.get("mask", kwargs.get("attn_mask"))

        B, N, D = hidden_states.shape
        H = self.num_heads
        d = self.head_dim
        kv_H = self.num_kv_heads

        if encoder_hidden_states is not None:
            raise NotImplementedError("Cross-attention is not supported in ASVDAttention.")

        # Extract rotary kwargs if supplied
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_ids is None:
            position_ids = kwargs.get("position_ids")

        # 1. Project Q, K, V using low-rank factorized linear layers
        q: torch.Tensor = self.query(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous()
        k: torch.Tensor = self.key(hidden_states).view(B, N, kv_H, d).transpose(1, 2).contiguous()
        v: torch.Tensor = self.value(hidden_states).view(B, N, kv_H, d).transpose(1, 2).contiguous()

        # 2. Apply Rotary Position Embedding (RoPE) before GQA expansion and masking
        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, position_ids)
        elif self.rotary_emb is not None:
            try:
                # Use query tensor q (full H heads) to infer sequence length, shapes, and device correctly
                rotary_out = (
                    self.rotary_emb(q, position_ids)
                    if position_ids is not None
                    else self.rotary_emb(q)
                )
                if isinstance(rotary_out, (tuple, list)) and len(rotary_out) >= 2:
                    cos, sin = rotary_out[0], rotary_out[1]
                    q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # 3. GQA Expansion: Expand key/value heads to match query heads if kv_H < H
        if kv_H < H:
            k = repeat_kv(k, H // kv_H)
            v = repeat_kv(v, H // kv_H)

        # 4. Attention Scoring strictly in FP32 to prevent FP16 overflow
        scale: float = 1.0 / math.sqrt(d)
        scores: torch.Tensor = (
            torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
        )

        # 5. Causal Masking for Decoder Models
        if getattr(self, "is_decoder", False):
            causal_mask = torch.triu(
                torch.full((N, N), -10000.0, device=scores.device, dtype=scores.dtype),
                diagonal=1,
            )
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

        # 6. Padding Attention Mask (Broadcast 2D/3D to 4D additive logits format)
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                if attention_mask.dtype == torch.bool:
                    extended_mask = torch.zeros_like(attention_mask, dtype=scores.dtype)
                    extended_mask = extended_mask.masked_fill(~attention_mask, -10000.0)
                elif torch.is_floating_point(attention_mask) and (attention_mask < 0).any():
                    extended_mask = attention_mask.to(dtype=scores.dtype)
                else:
                    extended_mask = (1.0 - attention_mask.to(dtype=scores.dtype)) * -10000.0
                attention_mask = extended_mask.unsqueeze(1).unsqueeze(2)
            elif attention_mask.dim() == 3:
                if attention_mask.dtype == torch.bool:
                    extended_mask = torch.zeros_like(attention_mask, dtype=scores.dtype)
                    extended_mask = extended_mask.masked_fill(~attention_mask, -10000.0)
                elif torch.is_floating_point(attention_mask) and (attention_mask < 0).any():
                    extended_mask = attention_mask.to(dtype=scores.dtype)
                else:
                    extended_mask = (1.0 - attention_mask.to(dtype=scores.dtype)) * -10000.0
                attention_mask = extended_mask.unsqueeze(1)
            scores = apply_attention_mask(scores, attention_mask)

        # 7. Softmax & NaN Protection
        attn: torch.Tensor = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        # 8. Safe Head Mask Handling
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask_tensor = head_mask.view(1, -1, 1, 1)
            elif head_mask.dim() == 2:
                head_mask_tensor = head_mask.view(head_mask.size(0), -1, 1, 1)
            elif head_mask.dim() == 3:
                head_mask_tensor = head_mask.unsqueeze(1)
            else:
                head_mask_tensor = head_mask
            attn = attn * head_mask_tensor

        # 9. Context Aggregation
        context: torch.Tensor = torch.matmul(attn, v)
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        # 10. Output Projection (Bypassed for nested encoders to prevent double projection)
        if getattr(self, "is_nested_encoder", False):
            out = context
        else:
            out = self.out_proj(context)

        return out, attn if output_attentions else None

Hardened ASVD Attention with Factored Low-Rank Matrix Decomposition, GQA, RoPE, and LoRA.

Computational Complexity: - Factorized Projections: $\mathcal{O}(B \cdot N \cdot r \cdot (d + d_{\text{all}}))$ - Attention Scoring: $\mathcal{O}(B \cdot H \cdot N^2 \cdot d_k)$

Tensor Dimensions: - hidden_states: $[B, N, D]$ where $B$ is batch size, $N$ is sequence length, and $D$ is hidden dimension. - Output: Tuple of $([B, N, D], \text{None})$.

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 triton_kernel : str

The type of the None singleton.

Instance variables

prop c_proj : ASVDLinear

Expand source code

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

prop dense : ASVDLinear

Expand source code

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

prop k_lin : ASVDLinear

Expand source code

@property
def k_lin(self) -> ASVDLinear:
    return self.key

prop k_proj : ASVDLinear

Expand source code

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

prop o_proj : ASVDLinear

Expand source code

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

prop out_lin : ASVDLinear

Expand source code

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

prop q_lin : ASVDLinear

Expand source code

@property
def q_lin(self) -> ASVDLinear:
    return self.query

prop q_proj : ASVDLinear

Expand source code

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

prop v_lin : ASVDLinear

Expand source code

@property
def v_lin(self) -> ASVDLinear:
    return self.value

prop v_proj : ASVDLinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    """Extracts and factorizes pretrained weights across DistilBERT, RoBERTa, GPT-2, Qwen/LLaMA."""
    extracted_weights: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

    # 1. Robust Nested Encoder Detection (prior to extraction)
    # Robustly detects both:
    # a) Outer RobertaAttention / BertAttention with inner 'self' (query/key/value) and 'output.dense'
    # b) Inner RobertaSelfAttention / BertSelfAttention
    # Prevents double output projection signal attenuation
    cls_name = src.__class__.__name__
    has_out_lin = hasattr(src, "out_lin")

    is_nested_encoder = False
    if (
        hasattr(src, "self")
        and hasattr(src.self, "query")
        and hasattr(src, "output")
        and hasattr(src.output, "dense")
    ):
        # Outer RobertaAttention containing inner self-attention and output dense
        is_nested_encoder = True
    elif (
        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:
        # Inner SelfAttention module
        is_nested_encoder = True

    self.is_nested_encoder = bool(is_nested_encoder)

    # 2. GPT-2: Fused Conv1D c_attn + Conv1D c_proj
    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()
            c_b = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            c_b = c_b.data.clone() if c_b is not None else None
            c_cls = getattr(base_c, "__class__", type(base_c)).__name__
            if (
                ("Conv1D" in c_cls)
                or hasattr(base_c, "nf")
                or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                or hasattr(c_attn_mod, "nf")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                c_w = c_w.t().contiguous()
            split_size = c_w.size(0) // 3
            extracted_weights["q"] = (
                c_w[:split_size, :],
                c_b[:split_size] if c_b is not None else None,
            )
            extracted_weights["k"] = (
                c_w[split_size : 2 * split_size, :],
                c_b[split_size : 2 * split_size] if c_b is not None else None,
            )
            extracted_weights["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()
                o_b = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                o_b = o_b.data.clone() if o_b is not None else None
                p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_o, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                ):
                    o_w = o_w.t().contiguous()
                extracted_weights["o"] = (o_w, o_b)

    # 3. RoBERTa / BERT Nested Encoder: self.query, self.key, self.value
    s = getattr(src, "self", src)
    name_map = {"q": "query", "k": "key", "v": "value"}
    for proj, attr in name_map.items():
        if (
            extracted_weights[proj][0] is None
            and hasattr(s, attr)
            and hasattr(getattr(s, attr), "weight")
        ):
            mod = getattr(s, attr)
            b = getattr(mod, "bias", None)
            extracted_weights[proj] = (
                mod.weight.data,
                b.data if b is not None else None,
            )

    if (
        not self.is_nested_encoder
        and extracted_weights["o"][0] is None
        and hasattr(src, "output")
        and hasattr(src.output, "dense")
    ):
        mod = src.output.dense
        b = getattr(mod, "bias", None)
        extracted_weights["o"] = (
            mod.weight.data,
            b.data if b is not None else None,
        )

    # 4. Qwen / LLaMA: q_proj, k_proj, v_proj, o_proj
    for proj in ["q", "k", "v", "o"]:
        if extracted_weights[proj][0] is None:
            for candidate in [f"{proj}_proj", f"{proj}_lin", proj]:
                if hasattr(src, candidate) and hasattr(getattr(src, candidate), "weight"):
                    mod = getattr(src, candidate)
                    b = getattr(mod, "bias", None)
                    extracted_weights[proj] = (
                        mod.weight.data,
                        b.data if b is not None else None,
                    )
                    break

    # 5. Fallback to UniversalAdapter
    for proj in ["q", "k", "v", "o"]:
        if extracted_weights[proj][0] is None:
            try:
                w, b = UniversalAdapter.get_weight_bias(src, proj)
                if w is not None:
                    extracted_weights[proj] = (w, b)
            except Exception as e:
                print(f"copy_weights failed: {e}")
                logger.warning(f"copy_weights failed: {e}")

    # 6. Check for missing essential projections and Frobenius norm verification
    missing = [p for p in ["q", "k", "v"] if extracted_weights[p][0] is None]
    if not self.is_nested_encoder and extracted_weights["o"][0] is None:
        missing.append("o")
    if missing:
        raise RuntimeError(
            f"ASVDAttention.copy_weights: could not extract mandatory projections {missing} "
            f"from {src.__class__.__name__}. Refusing to install randomly-initialised weights."
        )

    for p in ["q", "k", "v"]:
        w_check = extracted_weights[p][0]
        if w_check is not None:
            norm = float(torch.norm(w_check.float()).item())
            if norm < 1e-4:
                logger.warning(
                    f"ASVDAttention extracted weight for '{p}' has near-zero Frobenius norm ({norm:.2e})"
                )

    if extracted_weights["o"][0] is not None:
        norm_o = float(torch.norm(extracted_weights["o"][0].float()).item())
        if norm_o < 1e-4:
            logger.warning(
                f"ASVDAttention extracted weight for 'o' has near-zero Frobenius norm ({norm_o:.2e})"
            )

    # 7. SVD Factorization & Reconstruction Verification
    for proj, target in [
        ("q", self.query),
        ("k", self.key),
        ("v", self.value),
        ("o", self.out_proj),
    ]:
        if proj == "o" and self.is_nested_encoder:
            self.out_proj.set_identity()
            continue
        w, b = extracted_weights[proj]
        if w is not None:
            self._decompose_and_set(target, w, b, proj_name=proj)

    # 8. Inherit Decoder & Rotary Metadata
    if hasattr(src, "rotary_emb"):
        self.rotary_emb = src.rotary_emb
    if hasattr(src, "is_decoder"):
        self.is_decoder = bool(src.is_decoder)
    elif hasattr(src, "c_attn") or cls_name in (
        "GPT2Attention",
        "Qwen2Attention",
        "Qwen2_5_VLAttention",
        "QwenAttention",
    ):
        self.is_decoder = True

Extracts and factorizes pretrained weights across DistilBERT, RoBERTa, GPT-2, Qwen/LLaMA.

def forward(self,
hidden_states: torch.Tensor | None = None,
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_ids: torch.Tensor | None = None,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
*args: Any,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor | None]

Expand source code

def forward(
    self,
    hidden_states: torch.Tensor | None = None,
    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_ids: torch.Tensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Forward pass for factorized ASVD attention."""
    if hidden_states is None:
        hidden_states = kwargs.get("query", kwargs.get("x"))
    if hidden_states is None and len(args) > 0 and isinstance(args[0], torch.Tensor):
        hidden_states = args[0]
    if attention_mask is None:
        attention_mask = kwargs.get("mask", kwargs.get("attn_mask"))

    B, N, D = hidden_states.shape
    H = self.num_heads
    d = self.head_dim
    kv_H = self.num_kv_heads

    if encoder_hidden_states is not None:
        raise NotImplementedError("Cross-attention is not supported in ASVDAttention.")

    # Extract rotary kwargs if supplied
    if position_embeddings is None:
        position_embeddings = kwargs.get("position_embeddings")
    if position_ids is None:
        position_ids = kwargs.get("position_ids")

    # 1. Project Q, K, V using low-rank factorized linear layers
    q: torch.Tensor = self.query(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous()
    k: torch.Tensor = self.key(hidden_states).view(B, N, kv_H, d).transpose(1, 2).contiguous()
    v: torch.Tensor = self.value(hidden_states).view(B, N, kv_H, d).transpose(1, 2).contiguous()

    # 2. Apply Rotary Position Embedding (RoPE) before GQA expansion and masking
    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, position_ids)
    elif self.rotary_emb is not None:
        try:
            # Use query tensor q (full H heads) to infer sequence length, shapes, and device correctly
            rotary_out = (
                self.rotary_emb(q, position_ids)
                if position_ids is not None
                else self.rotary_emb(q)
            )
            if isinstance(rotary_out, (tuple, list)) and len(rotary_out) >= 2:
                cos, sin = rotary_out[0], rotary_out[1]
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        except Exception as e:
            logger.debug("UniversalAdapter fallback failed: %s", e)

    # 3. GQA Expansion: Expand key/value heads to match query heads if kv_H < H
    if kv_H < H:
        k = repeat_kv(k, H // kv_H)
        v = repeat_kv(v, H // kv_H)

    # 4. Attention Scoring strictly in FP32 to prevent FP16 overflow
    scale: float = 1.0 / math.sqrt(d)
    scores: torch.Tensor = (
        torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
    )

    # 5. Causal Masking for Decoder Models
    if getattr(self, "is_decoder", False):
        causal_mask = torch.triu(
            torch.full((N, N), -10000.0, device=scores.device, dtype=scores.dtype),
            diagonal=1,
        )
        scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

    # 6. Padding Attention Mask (Broadcast 2D/3D to 4D additive logits format)
    if attention_mask is not None:
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                extended_mask = torch.zeros_like(attention_mask, dtype=scores.dtype)
                extended_mask = extended_mask.masked_fill(~attention_mask, -10000.0)
            elif torch.is_floating_point(attention_mask) and (attention_mask < 0).any():
                extended_mask = attention_mask.to(dtype=scores.dtype)
            else:
                extended_mask = (1.0 - attention_mask.to(dtype=scores.dtype)) * -10000.0
            attention_mask = extended_mask.unsqueeze(1).unsqueeze(2)
        elif attention_mask.dim() == 3:
            if attention_mask.dtype == torch.bool:
                extended_mask = torch.zeros_like(attention_mask, dtype=scores.dtype)
                extended_mask = extended_mask.masked_fill(~attention_mask, -10000.0)
            elif torch.is_floating_point(attention_mask) and (attention_mask < 0).any():
                extended_mask = attention_mask.to(dtype=scores.dtype)
            else:
                extended_mask = (1.0 - attention_mask.to(dtype=scores.dtype)) * -10000.0
            attention_mask = extended_mask.unsqueeze(1)
        scores = apply_attention_mask(scores, attention_mask)

    # 7. Softmax & NaN Protection
    attn: torch.Tensor = F.softmax(scores, dim=-1).to(hidden_states.dtype)
    attn = torch.nan_to_num(attn, nan=0.0)
    attn = self.dropout(attn)

    # 8. Safe Head Mask Handling
    if head_mask is not None:
        if head_mask.dim() == 1:
            head_mask_tensor = head_mask.view(1, -1, 1, 1)
        elif head_mask.dim() == 2:
            head_mask_tensor = head_mask.view(head_mask.size(0), -1, 1, 1)
        elif head_mask.dim() == 3:
            head_mask_tensor = head_mask.unsqueeze(1)
        else:
            head_mask_tensor = head_mask
        attn = attn * head_mask_tensor

    # 9. Context Aggregation
    context: torch.Tensor = torch.matmul(attn, v)
    context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

    # 10. Output Projection (Bypassed for nested encoders to prevent double projection)
    if getattr(self, "is_nested_encoder", False):
        out = context
    else:
        out = self.out_proj(context)

    return out, attn if output_attentions else None

Forward pass for factorized ASVD attention.

Inherited members

• BaseAttention:
  • extra_repr

class ASVDLinear (in_features: int,
out_features: int,
rank: int,
bias: bool = True,
lora_rank: int = 0,
freeze_base: bool = False,
dtype: torch.dtype | None = None,
device: torch.device | None = None)

Expand source code

class ASVDLinear(nn.Module):
    is_custom_attention: bool = True
    r"""Low-rank factorized linear layer with explicit singular value scaling and trainable LoRA adapters:
    
    $$y = ((x V^\top) \odot \Sigma) U^\top + (x A^\top) B^\top + b$$
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        lora_rank: int = 0,
        freeze_base: bool = False,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.lora_rank = lora_rank
        self.freeze_base = freeze_base

        factory_kwargs = {"dtype": dtype, "device": device}
        self.U = nn.Parameter(torch.empty(out_features, rank, **factory_kwargs))
        self.S = nn.Parameter(torch.empty(rank, **factory_kwargs))
        self.V = nn.Parameter(torch.empty(rank, in_features, **factory_kwargs))

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

        # Trainable LoRA adapter parameters
        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))
            self.lora_scaling = 1.0
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)
            self.lora_scaling = 1.0

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

    def reset_parameters(self) -> None:
        with torch.no_grad():
            if self.U.dtype in (torch.float16, torch.bfloat16):
                u_fp32 = torch.empty_like(self.U, dtype=torch.float32, device="cpu")
                nn.init.orthogonal_(u_fp32)
                self.U.copy_(u_fp32.to(dtype=self.U.dtype, device=self.U.device))
            else:
                nn.init.orthogonal_(self.U)

            nn.init.ones_(self.S)

            if self.V.dtype in (torch.float16, torch.bfloat16):
                v_fp32 = torch.empty_like(self.V, dtype=torch.float32, device="cpu")
                nn.init.orthogonal_(v_fp32)
                self.V.copy_(v_fp32.to(dtype=self.V.dtype, device=self.V.device))
            else:
                nn.init.orthogonal_(self.V)

            if self.bias is not None:
                nn.init.zeros_(self.bias)

            if getattr(self, "lora_A", None) is not None 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 getattr(self, "lora_B", None) is not None 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:
        """Freeze U, S, V and bias parameters so only LoRA adapters adapt during recovery fine-tuning."""
        self.U.requires_grad = False
        self.S.requires_grad = False
        self.V.requires_grad = False
        if self.bias is not None:
            self.bias.requires_grad = False
        if getattr(self, "lora_A", None) is not None and self.lora_A is not None:
            self.lora_A.requires_grad = True
        if getattr(self, "lora_B", None) is not None and self.lora_B is not None:
            self.lora_B.requires_grad = True

    def unfreeze_base_weights(self) -> None:
        """Unfreeze base weights for standard end-to-end training."""
        self.U.requires_grad = True
        self.S.requires_grad = True
        self.V.requires_grad = True
        if self.bias is not None:
            self.bias.requires_grad = True

    def set_svd_factors(
        self,
        U: torch.Tensor,
        V: torch.Tensor,
        S: torch.Tensor | None = None,
        bias: torch.Tensor | None = None,
    ) -> None:
        with torch.no_grad():
            r = min(U.size(1), V.size(0))
            if r != self.rank:
                device, dtype = self.U.device, self.U.dtype
                req_grad = self.U.requires_grad
                self.rank = r
                self.U = nn.Parameter(
                    torch.zeros(self.out_features, r, device=device, dtype=dtype),
                    requires_grad=req_grad,
                )
                self.S = nn.Parameter(
                    torch.ones(r, device=device, dtype=dtype), requires_grad=req_grad
                )
                self.V = nn.Parameter(
                    torch.zeros(r, self.in_features, device=device, dtype=dtype),
                    requires_grad=req_grad,
                )

            self.U.data.zero_()
            self.V.data.zero_()
            self.S.data.zero_()

            out_copy = min(self.out_features, U.size(0))
            in_copy = min(self.in_features, V.size(1))

            self.U.data[:out_copy, :r].copy_(
                U[:out_copy, :r].to(device=self.U.device, dtype=self.U.dtype)
            )
            self.V.data[:r, :in_copy].copy_(
                V[:r, :in_copy].to(device=self.V.device, dtype=self.V.dtype)
            )
            if S is not None:
                self.S.data[:r].copy_(S[:r].to(device=self.S.device, dtype=self.S.dtype))
            else:
                self.S.data[:r].fill_(1.0)

            if bias is not None and self.bias is not None:
                min_b = min(self.out_features, bias.numel())
                self.bias.data[:min_b].copy_(
                    bias[:min_b].to(device=self.bias.device, dtype=self.bias.dtype)
                )
            elif self.bias is not None:
                self.bias.data.zero_()

    def set_identity(self) -> None:
        """Initializes factorized linear transformation to act as an identity pass-through."""
        with torch.no_grad():
            self.U.data.zero_()
            self.V.data.zero_()
            self.S.data.fill_(1.0)
            r = min(self.rank, self.in_features, self.out_features)
            self.U.data[:r, :r].fill_diagonal_(1.0)
            self.V.data[:r, :r].fill_diagonal_(1.0)
            if self.bias is not None:
                self.bias.data.zero_()
            self.freeze_base_weights()

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for inspection & compatibility."""
        W = torch.matmul(self.U * self.S.unsqueeze(0), self.V)
        if getattr(self, "lora_A", None) is not None and self.lora_A is not None:
            W = W + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
        return W

    triton_kernel: str = "low_rank"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (*, in_features)
        # Factorized SVD projection: y = ((x @ V.T) * S) @ U.T + bias
        V_eff = self.S.view(-1, 1) * self.V
        out = triton_dispatch(self.triton_kernel, x, self.U, V_eff, bias=self.bias)

        # 4. Residual LoRA adapter path: (x @ lora_A.T) @ lora_B.T
        if getattr(self, "lora_A", None) is not None and self.lora_A is not None:
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
            out = out + lora_out * self.lora_scaling

        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

Low-rank factorized linear layer with explicit singular value scaling and trainable LoRA adapters:

$$y = ((x V^\top) \odot \Sigma) U^\top + (x A^\top) B^\top + b$$

var triton_kernel : str

The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for inspection & compatibility."""
    W = torch.matmul(self.U * self.S.unsqueeze(0), self.V)
    if getattr(self, "lora_A", None) is not None and self.lora_A is not None:
        W = W + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
    return W

Reconstructed effective weight matrix for inspection & compatibility.

Methods

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: (*, in_features)
    # Factorized SVD projection: y = ((x @ V.T) * S) @ U.T + bias
    V_eff = self.S.view(-1, 1) * self.V
    out = triton_dispatch(self.triton_kernel, x, self.U, V_eff, bias=self.bias)

    # 4. Residual LoRA adapter path: (x @ lora_A.T) @ lora_B.T
    if getattr(self, "lora_A", None) is not None and self.lora_A is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
        out = out + lora_out * self.lora_scaling

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    """Freeze U, S, V and bias parameters so only LoRA adapters adapt during recovery fine-tuning."""
    self.U.requires_grad = False
    self.S.requires_grad = False
    self.V.requires_grad = False
    if self.bias is not None:
        self.bias.requires_grad = False
    if getattr(self, "lora_A", None) is not None and self.lora_A is not None:
        self.lora_A.requires_grad = True
    if getattr(self, "lora_B", None) is not None and self.lora_B is not None:
        self.lora_B.requires_grad = True

Freeze U, S, V and bias parameters so only LoRA adapters adapt during recovery fine-tuning.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    with torch.no_grad():
        if self.U.dtype in (torch.float16, torch.bfloat16):
            u_fp32 = torch.empty_like(self.U, dtype=torch.float32, device="cpu")
            nn.init.orthogonal_(u_fp32)
            self.U.copy_(u_fp32.to(dtype=self.U.dtype, device=self.U.device))
        else:
            nn.init.orthogonal_(self.U)

        nn.init.ones_(self.S)

        if self.V.dtype in (torch.float16, torch.bfloat16):
            v_fp32 = torch.empty_like(self.V, dtype=torch.float32, device="cpu")
            nn.init.orthogonal_(v_fp32)
            self.V.copy_(v_fp32.to(dtype=self.V.dtype, device=self.V.device))
        else:
            nn.init.orthogonal_(self.V)

        if self.bias is not None:
            nn.init.zeros_(self.bias)

        if getattr(self, "lora_A", None) is not None 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 getattr(self, "lora_B", None) is not None and self.lora_B is not None:
            nn.init.zeros_(self.lora_B)
            self.lora_B.requires_grad = True

def set_identity(self) ‑> None

Expand source code

def set_identity(self) -> None:
    """Initializes factorized linear transformation to act as an identity pass-through."""
    with torch.no_grad():
        self.U.data.zero_()
        self.V.data.zero_()
        self.S.data.fill_(1.0)
        r = min(self.rank, self.in_features, self.out_features)
        self.U.data[:r, :r].fill_diagonal_(1.0)
        self.V.data[:r, :r].fill_diagonal_(1.0)
        if self.bias is not None:
            self.bias.data.zero_()
        self.freeze_base_weights()

Initializes factorized linear transformation to act as an identity pass-through.

def set_svd_factors(self,
U: torch.Tensor,
V: torch.Tensor,
S: torch.Tensor | None = None,
bias: torch.Tensor | None = None) ‑> None

Expand source code

def set_svd_factors(
    self,
    U: torch.Tensor,
    V: torch.Tensor,
    S: torch.Tensor | None = None,
    bias: torch.Tensor | None = None,
) -> None:
    with torch.no_grad():
        r = min(U.size(1), V.size(0))
        if r != self.rank:
            device, dtype = self.U.device, self.U.dtype
            req_grad = self.U.requires_grad
            self.rank = r
            self.U = nn.Parameter(
                torch.zeros(self.out_features, r, device=device, dtype=dtype),
                requires_grad=req_grad,
            )
            self.S = nn.Parameter(
                torch.ones(r, device=device, dtype=dtype), requires_grad=req_grad
            )
            self.V = nn.Parameter(
                torch.zeros(r, self.in_features, device=device, dtype=dtype),
                requires_grad=req_grad,
            )

        self.U.data.zero_()
        self.V.data.zero_()
        self.S.data.zero_()

        out_copy = min(self.out_features, U.size(0))
        in_copy = min(self.in_features, V.size(1))

        self.U.data[:out_copy, :r].copy_(
            U[:out_copy, :r].to(device=self.U.device, dtype=self.U.dtype)
        )
        self.V.data[:r, :in_copy].copy_(
            V[:r, :in_copy].to(device=self.V.device, dtype=self.V.dtype)
        )
        if S is not None:
            self.S.data[:r].copy_(S[:r].to(device=self.S.device, dtype=self.S.dtype))
        else:
            self.S.data[:r].fill_(1.0)

        if bias is not None and self.bias is not None:
            min_b = min(self.out_features, bias.numel())
            self.bias.data[:min_b].copy_(
                bias[:min_b].to(device=self.bias.device, dtype=self.bias.dtype)
            )
        elif self.bias is not None:
            self.bias.data.zero_()

def unfreeze_base_weights(self) ‑> None

Expand source code

def unfreeze_base_weights(self) -> None:
    """Unfreeze base weights for standard end-to-end training."""
    self.U.requires_grad = True
    self.S.requires_grad = True
    self.V.requires_grad = True
    if self.bias is not None:
        self.bias.requires_grad = True

Unfreeze base weights for standard end-to-end training.

class AdaSplashAttention (config: Any,
rank: int | None = None,
alpha: float = 1.1,
n_iter: int = 15,
splash_radius: int = 128,
pretrained_attn: torch.nn.modules.module.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "adasplash",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
    supports_rope=True,
)
@ATTENTION_REGISTRY.register(
    "adasplash_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
    supports_rope=True,
)
class AdaSplashAttention(BaseAttention):
    r"""
    AdaSplash Attention Layer with α-Entmax and Local Splash Radius.
    """

    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,
        rank: int | None = None,
        alpha: float = 1.1,  # FIX: Softened from 1.5 to 1.1 for zero-shot accuracy
        n_iter: int = 15,
        splash_radius: int = 128,  # FIX: Increased from 32 to 128 for healthy context window
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank, **kwargs)
        if config is not None:
            self.adapter = ModelAdapter(config)
            self.alpha = float(getattr(config, "alpha", kwargs.get("alpha", alpha)))
            self.n_iter = getattr(config, "n_iter", kwargs.get("n_iter", n_iter))
            self.splash_radius = int(
                getattr(config, "splash_radius", kwargs.get("splash_radius", splash_radius))
            )
            lora_rank_cfg = getattr(
                config,
                "lora_rank",
                getattr(
                    config,
                    "residual_rank",
                    getattr(config, "rank", kwargs.get("lora_rank", kwargs.pop("rank", 16))),
                ),
            )
            self.lora_rank = int(lora_rank_cfg) if lora_rank_cfg is not None else 16
            self.num_heads = kwargs.get("num_heads", self.adapter.num_heads or 12)
            self.num_kv_heads = kwargs.get(
                "num_kv_heads", self.adapter.num_key_value_heads or self.num_heads
            )
            self.hidden_size = kwargs.get("hidden_size", self.adapter.hidden_size or 768)
            self.head_dim = kwargs.get(
                "head_dim",
                self.adapter.head_dim or (self.hidden_size // max(1, self.num_heads)),
            )
            self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
            dropout_p = getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                ),
            )
            self.is_decoder = getattr(config, "is_decoder", False)
            self.relative_attention_num_buckets = getattr(
                config, "relative_attention_num_buckets", 32
            )
            self.relative_attention_max_distance = getattr(
                config, "relative_attention_max_distance", 128
            )
        else:
            self.adapter = None
            self.alpha = float(kwargs.get("alpha", alpha))
            self.n_iter = int(kwargs.get("n_iter", n_iter))
            self.splash_radius = int(kwargs.get("splash_radius", splash_radius))
            self.lora_rank = int(kwargs.get("lora_rank", kwargs.pop("rank", 16)))
            self.hidden_size = kwargs.get("hidden_size", 768)
            self.num_heads = kwargs.get("num_heads", kwargs.get("num_attention_heads", 12))
            self.num_kv_heads = kwargs.get(
                "num_kv_heads", kwargs.get("num_key_value_heads", self.num_heads)
            )
            self.head_dim = kwargs.get("head_dim", self.hidden_size // max(1, self.num_heads))
            self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
            dropout_p = kwargs.get("dropout", 0.1)
            self.is_decoder = kwargs.get("is_decoder", False)
            self.relative_attention_num_buckets = 32
            self.relative_attention_max_distance = 128

        self.all_head_size = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim

        self.q_proj = AdaSplashLinear(
            self.hidden_size, self.all_head_size, lora_rank=self.lora_rank
        )
        self.k_proj = AdaSplashLinear(self.hidden_size, self.kv_head_size, lora_rank=self.lora_rank)
        self.v_proj = AdaSplashLinear(self.hidden_size, self.kv_head_size, lora_rank=self.lora_rank)
        self.out_proj = AdaSplashLinear(
            self.all_head_size, self.hidden_size, lora_rank=self.lora_rank
        )

        # Aliases
        self.q_lin = self.q_proj
        self.k_lin = self.k_proj
        self.v_lin = self.v_proj
        self.out_lin = self.out_proj
        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj
        self.dense = self.out_proj
        self.o_proj = self.out_proj
        self.c_proj = self.out_proj

        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

        self.is_nested_encoder = False
        self.has_relative_attention_bias = False
        self.relative_attention_bias = None

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

    def _relative_position_bucket(
        self,
        relative_position: torch.Tensor,
        bidirectional: bool = True,
        num_buckets: int = 32,
        max_distance: int = 128,
    ) -> torch.Tensor:
        relative_buckets = 0
        if bidirectional:
            num_buckets //= 2
            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
            relative_position = torch.abs(relative_position)
        else:
            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
        max_exact = num_buckets // 2
        is_small = relative_position < max_exact
        relative_position_if_large = max_exact + (
            torch.log(relative_position.float() / max_exact)
            / math.log(max_distance / max_exact)
            * (num_buckets - max_exact)
        ).to(torch.long)
        relative_position_if_large = torch.min(
            relative_position_if_large,
            torch.full_like(relative_position_if_large, num_buckets - 1),
        )
        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
        return relative_buckets

    def compute_bias(
        self, query_length: int, key_length: int, device: torch.device | None = None
    ) -> torch.Tensor:
        if self.relative_attention_bias is not None:
            context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
            memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
            relative_position = memory_position - context_position
            relative_position_bucket = self._relative_position_bucket(
                relative_position,
                bidirectional=(not self.is_decoder),
                num_buckets=self.relative_attention_num_buckets,
                max_distance=self.relative_attention_max_distance,
            )
            values = self.relative_attention_bias(relative_position_bucket)
            values = values.permute([2, 0, 1]).unsqueeze(0)
            return values
        return torch.zeros((1, self.num_heads, query_length, key_length), device=device)

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted_all = True
        extracted_weights = {}

        for proj, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            w, b = None, None
            try:
                w, b = UniversalAdapter.get_weight_bias(src, proj)
            except Exception:
                pass

            if w is None:
                # Output projection checks first
                if proj == "o":
                    if hasattr(src, "output") and hasattr(src.output, "dense"):
                        mod = src.output.dense
                        w = mod.weight
                        b = getattr(mod, "bias", None)
                    elif hasattr(src, "dense"):
                        mod = src.dense
                        w = mod.weight
                        b = getattr(mod, "bias", None)
                    elif hasattr(src, "c_proj"):
                        c_proj_mod = getattr(src.c_proj, "base_layer", src.c_proj)
                        w = c_proj_mod.weight.T
                        b = getattr(c_proj_mod, "bias", None)
                    elif hasattr(src, "out_lin"):
                        mod = getattr(src.out_lin, "base_layer", src.out_lin)
                        w = mod.weight
                        b = getattr(mod, "bias", None)
                    elif hasattr(src, "o_proj"):
                        mod = getattr(src.o_proj, "base_layer", src.o_proj)
                        w = mod.weight
                        b = getattr(mod, "bias", None)
                    elif hasattr(src, "wo"):
                        mod = getattr(src.wo, "base_layer", src.wo)
                        w = mod.weight
                        b = getattr(mod, "bias", None)
                # GPT-2 Fused c_attn
                elif hasattr(src, "c_attn") and proj in ["q", "k", "v"]:
                    c_attn_mod = getattr(src.c_attn, "base_layer", src.c_attn)
                    fused_w = c_attn_mod.weight.T
                    fused_b = getattr(c_attn_mod, "bias", None)
                    split_size = fused_w.size(0) // 3
                    idx = {"q": 0, "k": 1, "v": 2}[proj]
                    w = fused_w[idx * split_size : (idx + 1) * split_size, :]
                    b = (
                        fused_b[idx * split_size : (idx + 1) * split_size]
                        if fused_b is not None
                        else None
                    )
                # RoBERTa / BERT
                elif (
                    proj in ["q", "k", "v"] and hasattr(src, "self") and hasattr(src.self, "query")
                ):
                    s = src.self
                    name_map = {"q": "query", "k": "key", "v": "value"}
                    mod = getattr(s, name_map[proj])
                    w = mod.weight
                    b = getattr(mod, "bias", None)
                elif hasattr(src, "query") and proj == "q":
                    w = src.query.weight
                    b = getattr(src.query, "bias", None)
                elif hasattr(src, "key") and proj == "k":
                    w = src.key.weight
                    b = getattr(src.key, "bias", None)
                elif hasattr(src, "value") and proj == "v":
                    w = src.value.weight
                    b = getattr(src.value, "bias", None)
                # DistilBERT manual fallback
                elif hasattr(src, "q_lin") and proj == "q":
                    w = src.q_lin.weight
                    b = getattr(src.q_lin, "bias", None)
                elif hasattr(src, "k_lin") and proj == "k":
                    w = src.k_lin.weight
                    b = getattr(src.k_lin, "bias", None)
                elif hasattr(src, "v_lin") and proj == "v":
                    w = src.v_lin.weight
                    b = getattr(src.v_lin, "bias", None)
                elif hasattr(src, "out_lin") and proj == "o":
                    w = src.out_lin.weight
                    b = getattr(src.out_lin, "bias", None)
                # Modern Decoders
                elif hasattr(src, f"{proj}_proj"):
                    mod = getattr(src, f"{proj}_proj")
                    w = mod.weight
                    b = getattr(mod, "bias", None)

            extracted_weights[proj] = (w, b)
            if w is None:
                extracted_all = False

        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested_encoder = 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")
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            is_nested_encoder = True

        if not extracted_all:
            if (
                is_nested_encoder
                and extracted_weights["q"][0] is not None
                and extracted_weights["k"][0] is not None
                and extracted_weights["v"][0] is not None
            ):
                extracted_all = True

        if not extracted_all:
            raise RuntimeError(
                f"AdaSplashAttention.copy_weights: could not extract mandatory projections "
                f"{[p for p, (w, _) in extracted_weights.items() if w is None]} from {cls_name}. "
                f"Refusing to install randomly-initialised weights."
            )

        for proj, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            w, b = extracted_weights[proj]
            if w is not None:
                with torch.no_grad():
                    out_d, in_d = target.weight.shape
                    src_out, src_in = w.shape
                    target.weight.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                        w[: min(out_d, src_out), : min(in_d, src_in)]
                    )
                    if b is not None and target.bias is not None:
                        target.bias.data[: min(len(target.bias), len(b))].copy_(
                            b[: min(len(target.bias), len(b))]
                        )
                    if getattr(target, "lora_A", None) is not None:
                        nn.init.kaiming_uniform_(target.lora_A, a=math.sqrt(5))
                        target.lora_A.requires_grad = True
                    if getattr(target, "lora_B", None) is not None:
                        nn.init.zeros_(target.lora_B)
                        target.lora_B.requires_grad = True
            elif proj == "o" and is_nested_encoder:
                with torch.no_grad():
                    if target.weight.shape[0] == target.weight.shape[1]:
                        nn.init.eye_(target.weight)
                    else:
                        nn.init.zeros_(target.weight)
                        min_d = min(target.weight.shape[0], target.weight.shape[1])
                        target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        nn.init.zeros_(target.bias)
                    if getattr(target, "lora_A", None) is not None:
                        nn.init.kaiming_uniform_(target.lora_A, a=math.sqrt(5))
                        target.lora_A.requires_grad = True
                    if getattr(target, "lora_B", None) is not None:
                        nn.init.zeros_(target.lora_B)
                        target.lora_B.requires_grad = True

            target.weight.requires_grad = False
            if target.bias is not None:
                target.bias.requires_grad = False

        if hasattr(src, "has_relative_attention_bias"):
            self.has_relative_attention_bias = bool(src.has_relative_attention_bias)
        if hasattr(src, "relative_attention_bias") and src.relative_attention_bias is not None:
            self.relative_attention_bias = src.relative_attention_bias
            self.has_relative_attention_bias = True
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)

        self.is_nested_encoder = is_nested_encoder

    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, L_q, D = hidden_states.shape
        L_k = kv_input.shape[1]
        H, Dh = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        q = self.q_proj(hidden_states).view(B, L_q, H, Dh).transpose(1, 2).contiguous()
        k = self.k_proj(kv_input).view(B, L_k, self.num_kv_heads, Dh).transpose(1, 2).contiguous()
        v = self.v_proj(kv_input).view(B, L_k, self.num_kv_heads, Dh).transpose(1, 2).contiguous()

        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]

            def _rotate_half(x):
                _d = x.shape[-1]
                _h = _d // 2
                _rot = torch.cat((-x[..., _h : 2 * _h], x[..., :_h]), dim=-1)
                if _d % 2 == 1:
                    _rot = torch.cat([_rot, x[..., -1:]], dim=-1)
                return _rot

            if cos.dim() == 3:
                cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
            while cos.dim() < q.dim():
                cos, sin = cos.unsqueeze(0), sin.unsqueeze(0)
            rotary_dim = cos.shape[-1]
            cos_q, sin_q = (
                (cos[..., : q.shape[-2], :], sin[..., : q.shape[-2], :])
                if cos.shape[-2] != q.shape[-2]
                else (cos, sin)
            )
            cos_k, sin_k = (
                (cos[..., : k.shape[-2], :], sin[..., : k.shape[-2], :])
                if cos.shape[-2] != k.shape[-2]
                else (cos, sin)
            )
            if rotary_dim <= q.shape[-1]:
                q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
                k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
                q_rot = (q_rot * cos_q.to(q_rot.dtype)) + (
                    _rotate_half(q_rot) * sin_q.to(q_rot.dtype)
                )
                k_rot = (k_rot * cos_k.to(k_rot.dtype)) + (
                    _rotate_half(k_rot) * sin_k.to(k_rot.dtype)
                )
                q = torch.cat([q_rot, q_pass], dim=-1) if q_pass.numel() > 0 else q_rot
                k = torch.cat([k_rot, k_pass], dim=-1) if k_pass.numel() > 0 else k_rot
            else:
                cos_cut_q, sin_cut_q = (
                    cos_q[..., : q.shape[-1]],
                    sin_q[..., : q.shape[-1]],
                )
                cos_cut_k, sin_cut_k = (
                    cos_k[..., : k.shape[-1]],
                    sin_k[..., : k.shape[-1]],
                )
                q = (q * cos_cut_q.to(q.dtype)) + (_rotate_half(q) * sin_cut_q.to(q.dtype))
                k = (k * cos_cut_k.to(k.dtype)) + (_rotate_half(k) * sin_cut_k.to(k.dtype))
        elif hasattr(self, "rotary_emb") and self.rotary_emb is not None:
            try:
                q, k = maybe_apply_rotary(
                    q, k, self.adapter, self.rotary_emb, kwargs.get("position_ids")
                )
            except Exception:
                pass

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

        scale = Dh**-0.5
        logits = torch.matmul(q, k.transpose(-2, -1).contiguous()) * scale

        position_bias = kwargs.get("position_bias")
        if position_bias is None and self.has_relative_attention_bias:
            position_bias = self.compute_bias(L_q, L_k, device=device)

        mask = attention_mask if attention_mask is not None else kwargs.get("mask")
        if mask is not None:
            logits = apply_attention_mask(logits, mask)

        if position_bias is not None:
            logits = logits + position_bias.to(logits.dtype)

        mt = str(getattr(self.config, "model_type", "") or "").lower()
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_decoder = (
            self.is_decoder
            or getattr(self.config, "is_decoder", False)
            or any(mt.startswith(p) for p in decoder_prefixes)
        )
        causal_tril = None
        if is_decoder:
            causal_tril = (
                torch.tril(torch.ones((L_q, L_k), device=device, dtype=torch.bool))
                .unsqueeze(0)
                .unsqueeze(0)
            )
            logits = logits.masked_fill(~causal_tril, -10000.0)

        splash_mask_4d = None
        if (
            is_decoder
            and self.splash_radius > 0
            and not is_cross_attention
            and L_k > self.splash_radius
        ):
            q_idx = torch.arange(L_q, device=device).unsqueeze(1)
            k_idx = torch.arange(L_k, device=device).unsqueeze(0)
            dist = q_idx - k_idx
            splash_local = ((dist >= 0) & (dist <= self.splash_radius)) | (k_idx < 4)
            splash_mask_4d = splash_local.unsqueeze(0).unsqueeze(0)
            logits = logits.masked_fill(~splash_mask_4d, -10000.0)

        valid_mask_entmax = None
        if not is_cross_attention:
            valid_mask_entmax = logits > -5000.0

        if not is_cross_attention:
            attn = entmax_bisect(
                logits,
                alpha=self.alpha,
                valid_mask=valid_mask_entmax,
                dim=-1,
                n_iter=self.n_iter,
            ).to(dtype)
        else:
            attn = torch.nan_to_num(F.softmax(logits.float(), dim=-1), nan=0.0).to(dtype)

        attn = self.dropout(attn)

        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(-1, self.num_heads, 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn = attn * head_mask

        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, L_q, self.all_head_size)

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

        if position_bias is not None:
            return output, position_bias
        return output, (attn if output_attentions else None)

AdaSplash Attention Layer with α-Entmax and Local Splash Radius.

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.

Methods

def compute_bias(self, query_length: int, key_length: int, device: torch.device | None = None) ‑> torch.Tensor

Expand source code

def compute_bias(
    self, query_length: int, key_length: int, device: torch.device | None = None
) -> torch.Tensor:
    if self.relative_attention_bias is not None:
        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
        relative_position = memory_position - context_position
        relative_position_bucket = self._relative_position_bucket(
            relative_position,
            bidirectional=(not self.is_decoder),
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )
        values = self.relative_attention_bias(relative_position_bucket)
        values = values.permute([2, 0, 1]).unsqueeze(0)
        return values
    return torch.zeros((1, self.num_heads, query_length, key_length), device=device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AdaSplashLinear (in_features: int,
out_features: int,
lora_rank: int = 16,
bias: bool = True,
dtype: torch.dtype | None = None,
device: torch.device | None = None)

Expand source code

class AdaSplashLinear(nn.Module):
    r"""
    Linear projection layer with frozen base weights and trainable residual LoRA adapter.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        lora_rank: int = 16,
        bias: bool = True,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features: int = in_features
        self.out_features: int = out_features
        self.lora_rank: int = min(lora_rank, min(in_features, out_features))

        factory_kwargs = {"dtype": dtype, "device": device}

        # Base weights - frozen to maintain pretrained stability
        self.weight: nn.Parameter = nn.Parameter(
            torch.empty((out_features, in_features), **factory_kwargs),
            requires_grad=False,
        )

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

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

        self.reset_parameters()

    def reset_parameters(self) -> None:
        with torch.no_grad():
            if self.weight.dtype in (torch.float16, torch.bfloat16):
                w_fp32 = torch.empty_like(self.weight, dtype=torch.float32, device="cpu")
                nn.init.kaiming_uniform_(w_fp32, a=math.sqrt(5))
                self.weight.copy_(w_fp32.to(dtype=self.weight.dtype, device=self.weight.device))
            else:
                nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

            if self.bias is not None:
                nn.init.zeros_(self.bias)

            if (
                getattr(self, "lora_A", None) is not None
                and getattr(self, "lora_B", None) is not None
            ):
                if self.lora_B.dtype in (torch.float16, torch.bfloat16):
                    b_fp32 = torch.empty_like(self.lora_B, dtype=torch.float32, device="cpu")
                    nn.init.kaiming_uniform_(b_fp32, a=math.sqrt(5))
                    self.lora_B.copy_(b_fp32.to(dtype=self.lora_B.dtype, device=self.lora_B.device))
                else:
                    nn.init.kaiming_uniform_(self.lora_B, a=math.sqrt(5))
                nn.init.zeros_(self.lora_A)

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

Linear projection layer with frozen base weights and trainable residual LoRA adapter.

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 getattr(self, "lora_A", None) is not None and getattr(self, "lora_B", None) is not None:
        out = out + F.linear(F.linear(x, self.lora_A), self.lora_B)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    with torch.no_grad():
        if self.weight.dtype in (torch.float16, torch.bfloat16):
            w_fp32 = torch.empty_like(self.weight, dtype=torch.float32, device="cpu")
            nn.init.kaiming_uniform_(w_fp32, a=math.sqrt(5))
            self.weight.copy_(w_fp32.to(dtype=self.weight.dtype, device=self.weight.device))
        else:
            nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

        if self.bias is not None:
            nn.init.zeros_(self.bias)

        if (
            getattr(self, "lora_A", None) is not None
            and getattr(self, "lora_B", None) is not None
        ):
            if self.lora_B.dtype in (torch.float16, torch.bfloat16):
                b_fp32 = torch.empty_like(self.lora_B, dtype=torch.float32, device="cpu")
                nn.init.kaiming_uniform_(b_fp32, a=math.sqrt(5))
                self.lora_B.copy_(b_fp32.to(dtype=self.lora_B.dtype, device=self.lora_B.device))
            else:
                nn.init.kaiming_uniform_(self.lora_B, a=math.sqrt(5))
            nn.init.zeros_(self.lora_A)

class AdaptiveSVDAttention (config: Any,
svd_rank: int | None = None,
residual_rank: int = 16,
energy_threshold: float = 0.999,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "adaptive_svd_attn",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "adaptive_svd", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "adaptive_svd_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class AdaptiveSVDAttention(BaseAttention):
    """Multi-head attention with adaptive SVD + residual factorized projections."""

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

    def __init__(
        self,
        config: Any,
        svd_rank: int | None = None,
        residual_rank: int = 16,
        energy_threshold: float = 0.999,  # ASVD-FIX-02
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, **kwargs)

        # ASVD-FIX-12: always store config
        self.config = config

        # ===========================================================
        # ASVD-FIX-04: single source of truth for dims
        # ===========================================================
        def _first_attr(obj, names, default=None):
            if obj is None:
                return default
            for n in names:
                v = getattr(obj, n, None)
                if v is not None and v != 0:
                    return v
            return default

        # hidden_size
        hidden_size = (
            kwargs.get("hidden_size")
            or _first_attr(config, ("hidden_size", "d_model", "n_embd", "dim"))
            or _first_attr(pretrained_attn, ("hidden_size", "embed_dim", "all_head_size"))
        )
        if hidden_size is None:
            for attr in ("q_proj", "query", "q_lin", "c_attn"):
                mod = getattr(pretrained_attn, attr, None) if pretrained_attn else None
                if mod is not None and hasattr(mod, "in_features"):
                    hidden_size = mod.in_features
                    break
        self.hidden_size = int(hidden_size or 768)

        # num_heads
        num_heads = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or _first_attr(config, ("num_attention_heads", "num_heads", "n_heads", "n_head"))
            or _first_attr(pretrained_attn, ("num_heads", "num_attention_heads", "n_heads"))
        )
        self.num_heads = max(1, int(num_heads or 12))

        # num_kv_heads
        num_kv_heads = (
            kwargs.get("num_kv_heads")
            or kwargs.get("num_key_value_heads")
            or _first_attr(config, ("num_key_value_heads", "num_kv_heads"))
            or _first_attr(pretrained_attn, ("num_key_value_heads", "num_kv_heads"))
            or self.num_heads
        )
        self.num_kv_heads = max(1, min(int(num_kv_heads), self.num_heads))

        # head_dim
        head_dim = kwargs.get("head_dim") or _first_attr(config, ("head_dim", "d_kv"))
        if not head_dim and pretrained_attn is not None:
            for attr in ("q_proj", "query", "q_lin"):
                mod = getattr(pretrained_attn, attr, None)
                if mod is not None and hasattr(mod, "weight") and mod.weight is not None:
                    if mod.weight.ndim == 2:
                        head_dim = mod.weight.shape[0] // max(1, self.num_heads)
                        break
        self.head_dim = int(head_dim or (self.hidden_size // 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 // self.num_kv_heads)

        # ===========================================================
        # ASVD-FIX-08: is_causal computed once
        # ===========================================================
        _mt = str(getattr(config, "model_type", "") or "").lower() if config else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
            "codegen",
        )
        is_causal = (
            (is_decoder is True)
            or bool(kwargs.get("is_decoder", False))
            or bool(kwargs.get("is_causal", False))
            or bool(getattr(config, "is_decoder", False))
            if config
            else False
        )
        if not is_causal and config is not None:
            is_causal = (
                bool(getattr(config, "is_causal", False))
                or bool(getattr(config, "causal", False))
                or any(_mt.startswith(p) for p in _decoder_prefixes)
            )
        if not is_causal and pretrained_attn is not None:
            is_causal = (
                hasattr(pretrained_attn, "c_attn")
                or bool(getattr(pretrained_attn, "is_decoder", False))
                or bool(getattr(pretrained_attn, "is_causal", False))
            )
        self.is_decoder = bool(is_causal)
        self.is_causal = bool(is_causal)

        # ===========================================================
        # SVD configuration
        # ===========================================================
        cfg_svd_rank = getattr(config, "svd_rank", None) if config else None
        cfg_rank = getattr(config, "rank", None) if config else None
        resolved_svd_rank = svd_rank if svd_rank is not None else (cfg_svd_rank or cfg_rank or 64)
        # Clamp to legal range
        self.svd_rank = max(1, min(int(resolved_svd_rank), self.hidden_size, self.all_head_size))

        self.residual_rank = max(0, int(residual_rank))

        # ASVD-FIX-02: consistent threshold
        cfg_et = getattr(config, "energy_threshold", None) if config else None
        self.energy_threshold = float(cfg_et if cfg_et is not None else energy_threshold)
        # Clamp to valid range so rank genuinely varies with svd_rank
        self.energy_threshold = max(0.5, min(self.energy_threshold, 0.9999))

        dropout_p = 0.1
        if config is not None:
            dropout_p = getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                ),
            )
        self.dropout = nn.Dropout(float(dropout_p))

        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None) if pretrained_attn else None

        # Relative attention bias (T5)
        self.has_relative_attention_bias = False
        self.relative_attention_bias: nn.Embedding | None = None
        self.relative_attention_num_buckets = int(
            getattr(config, "relative_attention_num_buckets", 32) if config else 32
        )
        self.relative_attention_max_distance = int(
            getattr(config, "relative_attention_max_distance", 128) if config else 128
        )

        # ===========================================================
        # Projections
        # ===========================================================
        self.query = AdaptiveSVDLinear(
            self.hidden_size,
            self.all_head_size,
            self.svd_rank,
            self.residual_rank,
            bias=True,
        )
        self.key = AdaptiveSVDLinear(
            self.hidden_size,
            self.kv_head_size,
            self.svd_rank,
            self.residual_rank,
            bias=True,
        )
        self.value = AdaptiveSVDLinear(
            self.hidden_size,
            self.kv_head_size,
            self.svd_rank,
            self.residual_rank,
            bias=True,
        )
        self.out_proj = AdaptiveSVDLinear(
            self.all_head_size,
            self.hidden_size,
            self.svd_rank,
            self.residual_rank,
            bias=True,
        )

        # Aliases
        self.q_proj = self.query
        self.k_proj = self.key
        self.v_proj = self.value
        self.q_lin = self.query
        self.k_lin = self.key
        self.v_lin = self.value
        self.out_lin = self.out_proj
        self.dense = self.out_proj
        self.o_proj = self.out_proj
        self.c_proj = self.out_proj

        self.is_nested_encoder: bool = False

        # ASVD-FIX-11: storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

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

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

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

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

    # ------------------------------------------------------------------
    def _relative_position_bucket(
        self, relative_position, bidirectional=True, num_buckets=32, max_distance=128
    ):
        relative_buckets = 0
        if bidirectional:
            num_buckets //= 2
            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
            relative_position = torch.abs(relative_position)
        else:
            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
        max_exact = num_buckets // 2
        is_small = relative_position < max_exact
        relative_position_if_large = max_exact + (
            torch.log(relative_position.float() / max_exact)
            / math.log(max_distance / max_exact)
            * (num_buckets - max_exact)
        ).to(torch.long)
        relative_position_if_large = torch.min(
            relative_position_if_large,
            torch.full_like(relative_position_if_large, num_buckets - 1),
        )
        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
        return relative_buckets

    def compute_bias(self, query_length, key_length, device=None):
        if self.relative_attention_bias is None:
            return None
        ctx = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
        mem = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
        rel = mem - ctx
        rel_bucket = self._relative_position_bucket(
            rel,
            bidirectional=(not self.is_decoder),
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )
        values = self.relative_attention_bias(rel_bucket)
        return values.permute([2, 0, 1]).unsqueeze(0)

    # ------------------------------------------------------------------
    def _compute_energy_rank(self, S, max_rank, min_rank=4):
        """
        ASVD-FIX-02: with `energy_threshold` at 0.999, this returns a rank
        close to `max_rank`; with a lower threshold, more aggressive truncation.
        The actual return is always clamped to `[min_rank, max_rank]`.
        """
        s_sq = S**2
        total = s_sq.sum()
        if total <= 1e-12:
            return max(1, min(min_rank, max_rank))
        cum = torch.cumsum(s_sq, dim=0) / total
        idx = torch.where(cum >= self.energy_threshold)[0]
        r = int(idx[0].item()) + 1 if len(idx) > 0 else max_rank
        return max(1, min(max(r, min_rank), max_rank))

    def _init_svd_part(self, layer: AdaptiveSVDLinear, W_data, bias_data):
        with torch.no_grad():
            W = W_data.float()
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            max_r = min(layer.svd_rank, U.size(1), Vt.size(0))
            r = self._compute_energy_rank(S, max_rank=max_r)

            sqrt_s = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            U_r = (U[:, :r] * sqrt_s.unsqueeze(0)).to(layer.U.dtype)
            V_r = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(layer.V.dtype)

            layer.set_svd_factors(U_r, V_r)
            if layer.bias is not None and bias_data is not None:
                n = min(layer.bias.numel(), bias_data.numel())
                layer.bias.data[:n].copy_(bias_data[:n].to(layer.bias.dtype))

    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        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

        # ---- GPT-2 fused c_attn ----
        if hasattr(src, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            c = src.c_attn
            base_c = getattr(c, "base_layer", c)
            wt = getattr(base_c, "weight", getattr(c, "weight", None))
            if wt is not None:
                wt = wt.data.clone()
                is_conv = (
                    "Conv1D" in c.__class__.__name__
                    or "Conv1D" in base_c.__class__.__name__
                    or hasattr(c, "nf")
                    or hasattr(base_c, "nf")
                    or (wt.dim() == 2 and wt.shape[0] < wt.shape[1])
                )
                if is_conv or (wt.dim() == 2 and wt.shape[1] == 3 * self.hidden_size):
                    wt = wt.t().contiguous()
                b = getattr(base_c, "bias", getattr(c, "bias", None))
                bd = b.data.clone() if b is not None else None
                split = wt.shape[0] // 3
                extracted["q"] = (wt[:split], bd[:split] if bd is not None else None)
                extracted["k"] = (
                    wt[split : 2 * split],
                    bd[split : 2 * split] if bd is not None else None,
                )
                extracted["v"] = (
                    wt[2 * split :],
                    bd[2 * split :] if bd is not None else None,
                )

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            o = src.c_proj
            base_o = getattr(o, "base_layer", o)
            wo = getattr(base_o, "weight", getattr(o, "weight", None))
            if wo is not None:
                wo = wo.data.clone()
                is_conv_o = (
                    "Conv1D" in o.__class__.__name__
                    or "Conv1D" in base_o.__class__.__name__
                    or hasattr(o, "nf")
                    or hasattr(base_o, "nf")
                    or (
                        wo.dim() == 2
                        and wo.shape[0] != self.hidden_size
                        and wo.shape[1] == self.hidden_size
                    )
                    or (wo.dim() == 2 and wo.shape[0] < wo.shape[1])
                )
                if is_conv_o:
                    wo = wo.t().contiguous()
                bo = getattr(base_o, "bias", getattr(o, "bias", None))
                extracted["o"] = (wo, bo.data.clone() if bo is not None else None)

        # ---- BERT / RoBERTa nested ----
        if hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                if extracted[key][0] is None and hasattr(s, sub):
                    m = getattr(s, sub)
                    b = getattr(m, "bias", None)
                    extracted[key] = (m.weight.data, b.data if b is not None else None)

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            d = getattr(src.output.dense, "base_layer", src.output.dense)
            b = getattr(d, "bias", None)
            extracted["o"] = (d.weight.data, b.data if b is not None else None)

        # ---- flat modules ----
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        for key, names in named_map.items():
            if extracted[key][0] is None:
                for n in names:
                    if hasattr(src, n):
                        m = getattr(src, n)
                        if hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data,
                                b.data if b is not None else None,
                            )
                            break

        # ---- nested encoder detection ----
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, a) for a in ("dense", "out_proj", "c_proj", "out_lin", "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 not has_out_proj)
            )
            and not has_out_proj
        ):
            is_nested = True
        self.is_nested_encoder = is_nested

        # ---- tracking original bytes ----
        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # ---- validation ----
        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:
            raise RuntimeError(
                f"AdaptiveSVDAttention.copy_weights: missing {missing} from {cls_name}."
            )

        for k_ in ("q", "k", "v"):
            w = extracted[k_][0]
            if w is not None:
                n = torch.linalg.norm(w.float()).item()
                if n < 1e-4:
                    logger.warning("AdaptiveSVDAttention: %s near-zero norm (%.2e)", k_, n)

        # ---- factorize ----
        with torch.no_grad():
            for key, target in (
                ("q", self.query),
                ("k", self.key),
                ("v", self.value),
                ("o", self.out_proj),
            ):
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        # identity
                        eye = torch.eye(
                            target.out_features,
                            target.in_features,
                            device=target.U.device,
                            dtype=target.U.dtype,
                        )
                        self._init_svd_part(target, eye, None)
                    continue
                self._init_svd_part(target, w, b)

                # ASVD-FIX-10: freeze all base params + bias
                target.U.requires_grad = False
                target.V.requires_grad = False
                if target.bias is not None:
                    target.bias.requires_grad = False
                if target.lora_A is not None:
                    target.lora_A.requires_grad = True
                if target.lora_B is not None:
                    target.lora_B.requires_grad = True

        # ---- track compressed bytes ----
        try:
            for m in (self.query, self.key, self.value, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # ---- inherit auxiliary state from source ----
        if hasattr(src, "has_relative_attention_bias"):
            self.has_relative_attention_bias = bool(src.has_relative_attention_bias)
        if hasattr(src, "relative_attention_bias") and src.relative_attention_bias is not None:
            self.relative_attention_bias = src.relative_attention_bias
            self.has_relative_attention_bias = True
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)
            self.is_causal = self.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]:
        # ASVD-FIX-06: always return (output, attn_weights)
        key_value_states = kwargs.get("key_value_states", encoder_hidden_states)
        is_cross = key_value_states is not None
        kv_input = key_value_states if is_cross else hidden_states

        B, N_q, _ = hidden_states.shape
        N_k = kv_input.shape[1]
        H, D = self.num_heads, self.head_dim

        Q = self.query(hidden_states).view(B, N_q, H, D).transpose(1, 2).contiguous()
        K = self.key(kv_input).view(B, N_k, self.num_kv_heads, D).transpose(1, 2).contiguous()
        V = self.value(kv_input).view(B, N_k, self.num_kv_heads, D).transpose(1, 2).contiguous()

        # ---- RoPE (ASVD-FIX-05: uses module-level helper) ----
        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]
            if cos.dim() == 3:
                cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
            while cos.dim() < Q.dim():
                cos, sin = cos.unsqueeze(0), sin.unsqueeze(0)
            if cos.shape[-2] != Q.shape[-2]:
                cos, sin = cos[..., : Q.shape[-2], :], sin[..., : Q.shape[-2], :]
            rotary_dim = cos.shape[-1]
            if rotary_dim <= Q.shape[-1]:
                q_rot, q_pass = Q[..., :rotary_dim], Q[..., rotary_dim:]
                k_rot, k_pass = K[..., :rotary_dim], K[..., rotary_dim:]
                q_rot = (q_rot * cos.to(q_rot.dtype)) + (_rotate_half(q_rot) * sin.to(q_rot.dtype))
                k_rot = (k_rot * cos.to(k_rot.dtype)) + (_rotate_half(k_rot) * sin.to(k_rot.dtype))
                Q = torch.cat([q_rot, q_pass], dim=-1) if q_pass.numel() > 0 else q_rot
                K = torch.cat([k_rot, k_pass], dim=-1) if k_pass.numel() > 0 else k_rot
            else:
                cos_cut = cos[..., : Q.shape[-1]].to(Q.dtype)
                sin_cut = sin[..., : Q.shape[-1]].to(Q.dtype)
                Q = (Q * cos_cut) + (_rotate_half(Q) * sin_cut)
                cos_k = cos[..., : K.shape[-1]].to(K.dtype)
                sin_k = sin[..., : K.shape[-1]].to(K.dtype)
                K = (K * cos_k) + (_rotate_half(K) * sin_k)

        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = D**-0.5
        scores = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale

        # ---- Relative attention bias (T5) ----
        position_bias = kwargs.get("position_bias")
        if position_bias is None and self.has_relative_attention_bias:
            position_bias = self.compute_bias(N_q, N_k, device=hidden_states.device)

        mask = attention_mask if attention_mask is not None else kwargs.get("mask")

        if position_bias is not None:
            if mask is not None and mask.dim() == 4:
                position_bias = position_bias + mask.to(scores.dtype)
            scores = scores + position_bias.to(scores.dtype)
        elif mask is not None:
            scores = apply_attention_mask(scores, mask)

        # ---- Causal mask (ASVD-FIX-08: use self.is_causal) ----
        if self.is_causal and not is_cross and N_q > 1:
            causal_diag = 1 + max(0, N_k - N_q)
            causal = (
                torch.triu(
                    torch.full(
                        (N_q, N_k),
                        float("-inf"),
                        device=scores.device,
                        dtype=scores.dtype,
                    ),
                    diagonal=causal_diag,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores.float(), dim=-1).to(hidden_states.dtype)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights)

        if head_mask is not None:
            hm = head_mask.to(attn_weights.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            attn_weights = attn_weights * hm

        context = torch.matmul(attn_weights, V.to(attn_weights.dtype))
        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)

        output = output.to(hidden_states.dtype)
        return output, (attn_weights if output_attentions else None)

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("query", self.query),
            ("key", self.key),
            ("value", self.value),
            ("out_proj", self.out_proj),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["total"] = sum(out.values())
        return out

Multi-head attention with adaptive SVD + residual factorized projections.

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 k_lin_alias

Expand source code

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

prop q_lin_alias

Expand source code

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

prop v_lin_alias

Expand source code

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

Methods

def compute_bias(self, query_length, key_length, device=None)

Expand source code

def compute_bias(self, query_length, key_length, device=None):
    if self.relative_attention_bias is None:
        return None
    ctx = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
    mem = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
    rel = mem - ctx
    rel_bucket = self._relative_position_bucket(
        rel,
        bidirectional=(not self.is_decoder),
        num_buckets=self.relative_attention_num_buckets,
        max_distance=self.relative_attention_max_distance,
    )
    values = self.relative_attention_bias(rel_bucket)
    return values.permute([2, 0, 1]).unsqueeze(0)

def storage_bytes(self) ‑> dict

Expand source code

def storage_bytes(self) -> dict:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("query", self.query),
        ("key", self.key),
        ("value", self.value),
        ("out_proj", self.out_proj),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AdaptiveSVDLinear (in_features: int,
out_features: int,
svd_rank: int,
residual_rank: int = 16,
bias: bool = True,
trainable_svd: bool = False,
dtype: torch.dtype | None = None,
device: torch.device | None = None)

Expand source code

class AdaptiveSVDLinear(nn.Module):
    """
    Frozen SVD base + trainable LoRA residual.

    Forward:
        y = x V^T U^T + x lora_A^T lora_B^T + b

    Storage:
        U (out, r), V (r, in), lora_A (residual_rank, in), lora_B (out, residual_rank)
    """

    # ASVD-FIX-03: not attention
    is_custom_attention: bool = False
    is_svd_linear: bool = True
    is_compressed_svd: bool = True

    def __init__(
        self,
        in_features: int,
        out_features: int,
        svd_rank: int,
        residual_rank: int = 16,
        bias: bool = True,
        trainable_svd: bool = False,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        max_rank = min(self.in_features, self.out_features)
        self.svd_rank = int(min(svd_rank, max_rank))
        self.residual_rank = int(min(residual_rank, max_rank))

        fk = {"dtype": dtype, "device": device}
        self.U = nn.Parameter(
            torch.empty(self.out_features, self.svd_rank, **fk),
            requires_grad=trainable_svd,
        )
        self.V = nn.Parameter(
            torch.empty(self.svd_rank, self.in_features, **fk),
            requires_grad=trainable_svd,
        )

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

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

        self.reset_parameters()

    def reset_parameters(self) -> None:
        with torch.no_grad():
            for p in (self.U, self.V):
                if p.dtype in (torch.float16, torch.bfloat16):
                    tmp = torch.empty_like(p, dtype=torch.float32, device="cpu")
                    nn.init.orthogonal_(tmp)
                    p.copy_(tmp.to(dtype=p.dtype, device=p.device))
                else:
                    nn.init.orthogonal_(p)
            if self.lora_A is not None and self.lora_B is not None:
                nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
                nn.init.zeros_(self.lora_B)
            if self.bias is not None:
                nn.init.zeros_(self.bias)

    def set_svd_factors(self, U: torch.Tensor, V: torch.Tensor) -> None:
        with torch.no_grad():
            r = min(U.size(1), V.size(0), self.svd_rank)
            out_d = min(self.U.size(0), U.size(0))
            in_d = min(self.V.size(1), V.size(1))
            self.U.zero_()
            self.V.zero_()
            self.U[:out_d, :r].copy_(U[:out_d, :r].to(self.U.dtype))
            self.V[:r, :in_d].copy_(V[:r, :in_d].to(self.V.dtype))
            if self.lora_B is not None:
                nn.init.zeros_(self.lora_B)

    # ---- aliases for framework compatibility ----
    @property
    def A(self) -> nn.Parameter | None:
        return self.lora_A

    @property
    def B(self) -> nn.Parameter | None:
        return self.lora_B

    @property
    def weight(self) -> torch.Tensor:
        W = torch.matmul(self.U, self.V)
        if self.lora_A is not None and self.lora_B is not None:
            W = W + torch.matmul(self.lora_B, self.lora_A)
        return W

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # ASVD-FIX-09: dtype-safe
        orig_dtype = x.dtype
        V = self.V.to(x.dtype) if self.V.dtype != x.dtype else self.V
        U = self.U.to(x.dtype) if self.U.dtype != x.dtype else self.U

        out = F.linear(F.linear(x, V), U)

        if self.lora_A is not None and self.lora_B is not None:
            a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
            b = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
            out = out + F.linear(F.linear(x, a), b)

        if self.bias is not None:
            out = out + self.bias.to(out.dtype)

        return out.to(orig_dtype)

Frozen SVD base + trainable LoRA residual.

Forward -----= y = x V^T U^T + x lora_A^T lora_B^T + b

Storage -----= U (out, r), V (r, in), lora_A (residual_rank, in), lora_B (out, residual_rank)

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

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_compressed_svd : bool

The type of the None singleton.

var is_custom_attention : bool

The type of the None singleton.

var is_svd_linear : bool

The type of the None singleton.

Instance variables

prop A : torch.nn.parameter.Parameter | None

Expand source code

@property
def A(self) -> nn.Parameter | None:
    return self.lora_A

prop B : torch.nn.parameter.Parameter | None

Expand source code

@property
def B(self) -> nn.Parameter | None:
    return self.lora_B

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    W = torch.matmul(self.U, self.V)
    if self.lora_A is not None and self.lora_B is not None:
        W = W + torch.matmul(self.lora_B, self.lora_A)
    return W

Methods

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # ASVD-FIX-09: dtype-safe
    orig_dtype = x.dtype
    V = self.V.to(x.dtype) if self.V.dtype != x.dtype else self.V
    U = self.U.to(x.dtype) if self.U.dtype != x.dtype else self.U

    out = F.linear(F.linear(x, V), U)

    if self.lora_A is not None and self.lora_B is not None:
        a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
        b = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
        out = out + F.linear(F.linear(x, a), b)

    if self.bias is not None:
        out = out + self.bias.to(out.dtype)

    return out.to(orig_dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    with torch.no_grad():
        for p in (self.U, self.V):
            if p.dtype in (torch.float16, torch.bfloat16):
                tmp = torch.empty_like(p, dtype=torch.float32, device="cpu")
                nn.init.orthogonal_(tmp)
                p.copy_(tmp.to(dtype=p.dtype, device=p.device))
            else:
                nn.init.orthogonal_(p)
        if self.lora_A is not None and self.lora_B is not None:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

def set_svd_factors(self, U: torch.Tensor, V: torch.Tensor) ‑> None

Expand source code

def set_svd_factors(self, U: torch.Tensor, V: torch.Tensor) -> None:
    with torch.no_grad():
        r = min(U.size(1), V.size(0), self.svd_rank)
        out_d = min(self.U.size(0), U.size(0))
        in_d = min(self.V.size(1), V.size(1))
        self.U.zero_()
        self.V.zero_()
        self.U[:out_d, :r].copy_(U[:out_d, :r].to(self.U.dtype))
        self.V[:r, :in_d].copy_(V[:r, :in_d].to(self.V.dtype))
        if self.lora_B is not None:
            nn.init.zeros_(self.lora_B)

class AdaptiveSpanAttention (config: Any | None = None,
rank: int | None = None,
dropout: float = 0.1,
max_span: int = 128,
ramp_size: float = 32.0,
lora_rank: int = 16,
comp_config: Any | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
layer_idx: int | None = None,
return_style: str | None = None,
rolling_cache_size: int | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("adaptive_span")
@ATTENTION_REGISTRY.register("adaptive_span_attention")
class AdaptiveSpanAttention(BaseAttention):
    supports_causal: bool = True
    supports_bidirectional: bool = False
    supports_gqa: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        dropout: float = 0.1,
        max_span: int = 128,
        ramp_size: float = 32.0,
        lora_rank: int = 16,
        comp_config: Any | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        layer_idx: int | None = None,
        return_style: str | None = None,
        rolling_cache_size: int | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        # ---------------- Dimensions ----------------
        hidden_size = kwargs.get("d_model") or kwargs.get("hidden_size") or kwargs.get("dim")
        num_heads = (
            kwargs.get("num_heads") or kwargs.get("n_heads") or kwargs.get("num_attention_heads")
        )
        num_kv_heads = kwargs.get("num_kv_heads") or kwargs.get("num_key_value_heads")
        head_dim = 0

        if config is not None:
            adapter = ModelAdapter(config)
            hidden_size = hidden_size or adapter.hidden_size
            num_heads = num_heads or adapter.num_heads
            num_kv_heads = num_kv_heads or adapter.num_key_value_heads
            head_dim = getattr(adapter, "head_dim", 0)

        if not hidden_size and source_attn is not None:
            cand = getattr(
                source_attn,
                "hidden_size",
                getattr(
                    source_attn,
                    "embed_dim",
                    getattr(source_attn, "all_head_size", None),
                ),
            )
            if cand is None and hasattr(source_attn, "q_proj"):
                cand = getattr(source_attn.q_proj, "in_features", None)
            hidden_size = cand

        hidden_size = int(hidden_size or 768)
        num_heads = int(num_heads or 12)
        num_kv_heads = int(num_kv_heads or num_heads)
        head_dim = head_dim if head_dim > 0 else (hidden_size // num_heads if num_heads > 0 else 64)

        self.hidden_size: int = hidden_size
        self.d_model: int = hidden_size
        self.num_heads: int = num_heads
        self.num_kv_heads: int = num_kv_heads
        self.head_dim: int = 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 = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.layer_idx: int = int(layer_idx) if layer_idx is not None else 0
        self.rolling_cache_size: int | None = rolling_cache_size

        # ---------------- Span config ----------------
        if config is not None:
            self.max_span = int(getattr(config, "max_span", max_span))
            self.ramp_size = float(getattr(config, "ramp_size", ramp_size))
        elif comp_config is not None:
            self.max_span = int(getattr(comp_config, "max_span", max_span))
            self.ramp_size = float(getattr(comp_config, "ramp_size", ramp_size))
        else:
            self.max_span = int(max_span)
            self.ramp_size = float(ramp_size)

        # ---------------- Decoder / nested encoder detection ----------------
        model_type = getattr(config, "model_type", "").lower() if config else ""
        is_dec = (
            getattr(
                config,
                "is_decoder",
                getattr(config, "is_causal", kwargs.get("is_causal", False)),
            )
            if config
            else False
        )
        _dec_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if any(model_type.startswith(p) for p in _dec_prefixes):
            is_dec = True
        if hasattr(config, "architectures") and any(
            "GPT" in str(a) or "CausalLM" in str(a)
            for a in (getattr(config, "architectures", []) or [])
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)

        source_cls = source_attn.__class__.__name__ if source_attn is not None else ""
        cfg_model_type = getattr(config, "model_type", "").lower() if config else ""
        is_nested = (
            source_cls in ("BertSelfAttention", "RobertaSelfAttention")
            or cfg_model_type in ("roberta", "bert")
            or (
                source_cls.endswith(("SelfAttention", "SelfAttn"))
                and "distil" not in source_cls.lower()
                and "gpt" not in source_cls.lower()
                and "qwen" not in source_cls.lower()
            )
        )
        self.is_nested_encoder: bool = bool(kwargs.get("is_nested_encoder", is_nested))

        if return_style is None:
            if self.is_nested_encoder:
                return_style = "bert"
            elif "GPT2" in source_cls or "gpt2" in cfg_model_type:
                return_style = "gpt2"
            elif "T5" in source_cls or "t5" in cfg_model_type:
                return_style = "t5"
            elif self.is_decoder:
                return_style = "llama"
            else:
                return_style = "bert"
        self.return_style: str = str(return_style)

        self.has_relative_attention_bias: bool = False
        self.relative_attention_bias: nn.Embedding | None = None
        self.relative_attention_num_buckets: int = (
            getattr(config, "relative_attention_num_buckets", 32) if config else 32
        )
        self.relative_attention_max_distance: int = (
            getattr(config, "relative_attention_max_distance", 128) if config else 128
        )

        # ---------------- LoRA config ----------------
        if config is not None:
            resolved_lora_rank = getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", lora_rank)),
            )
            resolved_lora_alpha = getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        elif comp_config is not None:
            resolved_lora_rank = getattr(
                comp_config, "lora_rank", kwargs.get("lora_rank", lora_rank)
            )
            resolved_lora_alpha = getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
        else:
            resolved_lora_rank = kwargs.get("lora_rank", lora_rank)
            resolved_lora_alpha = kwargs.get("lora_alpha", 32.0)

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and source_attn is not None:
            for p in source_attn.parameters():
                target_dtype = p.dtype
                break

        self.lora_rank = int(resolved_lora_rank)
        self.lora_alpha = float(resolved_lora_alpha)
        self.freeze_base: bool = self.lora_rank > 0

        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=self.freeze_base,
            dtype=target_dtype,
        )
        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=self.freeze_base,
            dtype=target_dtype,
        )
        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=self.freeze_base,
            dtype=target_dtype,
        )
        if self.is_nested_encoder:
            self.out_proj: nn.Module = IdentityOutProj()
        else:
            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=self.freeze_base,
                dtype=target_dtype,
            )

        # =========================================================================
        # FIX: Initialized to 2.0 to prevent gradient vanishing
        # =========================================================================
        self.span_logits = nn.Parameter(torch.full((self.num_heads,), 2.0, dtype=torch.float32))

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", dropout),
                ),
            )
            if config
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)

        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )

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

        if target_dtype is not None:
            self.q_proj.to(target_dtype)
            self.k_proj.to(target_dtype)
            self.v_proj.to(target_dtype)
            self.out_proj.to(target_dtype)

    @property
    def current_span(self) -> torch.Tensor:
        return torch.sigmoid(self.span_logits) * self.max_span

    def get_span_regularization_loss(self, lambda_span: float = 0.01) -> torch.Tensor:
        return lambda_span * torch.sigmoid(self.span_logits).mean()

    def _relative_position_bucket(
        self, relative_position, bidirectional=True, num_buckets=32, max_distance=128
    ):
        relative_buckets = 0
        if bidirectional:
            num_buckets //= 2
            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
            relative_position = torch.abs(relative_position)
        else:
            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
        max_exact = num_buckets // 2
        is_small = relative_position < max_exact
        rel_large = max_exact + (
            torch.log(relative_position.float() / max_exact)
            / math.log(max_distance / max_exact)
            * (num_buckets - max_exact)
        ).to(torch.long)
        rel_large = torch.min(rel_large, torch.full_like(rel_large, num_buckets - 1))
        relative_buckets += torch.where(is_small, relative_position, rel_large)
        return relative_buckets

    def compute_bias(self, query_length: int, key_length: int, device=None) -> torch.Tensor:
        if self.relative_attention_bias is None:
            return torch.zeros((1, self.num_heads, query_length, key_length), device=device)
        ctx = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
        mem = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
        rel = mem - ctx
        buckets = self._relative_position_bucket(
            rel,
            bidirectional=(not self.is_decoder),
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )
        values = self.relative_attention_bias(buckets)
        return values.permute([2, 0, 1]).unsqueeze(0)

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        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 as e:
                logger.debug("UniversalAdapter fallback for %s: %s", proj, e)

        if hasattr(src, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            c_attn_layer = src.c_attn
            base_c = getattr(c_attn_layer, "base_layer", c_attn_layer)
            weight = getattr(base_c, "weight", getattr(c_attn_layer, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_layer, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                is_conv1d = (
                    hasattr(c_attn_layer, "nf")
                    or hasattr(base_c, "nf")
                    or c_attn_layer.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                )
                if is_conv1d and weight.dim() == 2:
                    weight = weight.t().contiguous()
                split = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split].contiguous(),
                    bias_data[:split].contiguous() if bias_data is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split].contiguous(),
                    (bias_data[split : 2 * split].contiguous() if bias_data is not None else None),
                )
                extracted["v"] = (
                    weight[2 * split :].contiguous(),
                    (bias_data[2 * split :].contiguous() if bias_data is not None else None),
                )

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            c_proj_layer = src.c_proj
            base_o = getattr(c_proj_layer, "base_layer", c_proj_layer)
            w_o = getattr(base_o, "weight", getattr(c_proj_layer, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                is_conv1d = (
                    hasattr(c_proj_layer, "nf")
                    or hasattr(base_o, "nf")
                    or c_proj_layer.__class__.__name__ == "Conv1D"
                    or base_o.__class__.__name__ == "Conv1D"
                )
                if is_conv1d and w_o.dim() == 2:
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj_layer, "bias", None))
                extracted["o"] = (
                    w_o.contiguous(),
                    b_o.data.clone().contiguous() if b_o is not None else None,
                )

        if hasattr(src, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = src.self
            if hasattr(s, "query") and extracted["q"][0] is None:
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key") and extracted["k"][0] is None:
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value") and extracted["v"][0] is None:
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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

        cls_name = src.__class__.__name__
        cfg_type = getattr(self.config, "model_type", "").lower() if self.config else ""
        if not self.is_nested_encoder:
            self.is_nested_encoder = bool(
                cls_name in ("BertSelfAttention", "RobertaSelfAttention")
                or cfg_type in ("roberta", "bert")
                or (
                    cls_name.endswith("SelfAttention")
                    and "distil" not in cls_name.lower()
                    and "gpt" not in cls_name.lower()
                    and "qwen" not in cls_name.lower()
                )
            )
        if self.is_nested_encoder and not isinstance(self.out_proj, IdentityOutProj):
            self.out_proj = IdentityOutProj()

        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("AdaptiveSpanAttention missing weights for projections: %s", missing)

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                if proj_key == "o" and self.is_nested_encoder:
                    continue
                w, b = extracted[proj_key]
                if w is not None:
                    out_d, in_d = target.weight.shape
                    src_out, src_in = w.shape
                    c_o, c_i = min(out_d, src_out), min(in_d, src_in)
                    target.weight.data[:c_o, :c_i].copy_(w[:c_o, :c_i])
                    if b is not None and target.bias is not None:
                        b_val = b.data if isinstance(b, nn.Parameter) else b
                        c_b = min(len(target.bias), len(b_val))
                        target.bias.data[:c_b].copy_(b_val[:c_b])
                target.freeze_base_weights()

        if getattr(src, "rotary_emb", None) is not None:
            self.rotary_emb = src.rotary_emb
        if hasattr(src, "has_relative_attention_bias"):
            self.has_relative_attention_bias = bool(src.has_relative_attention_bias)
        if getattr(src, "relative_attention_bias", None) is not None:
            self.relative_attention_bias = src.relative_attention_bias
            self.has_relative_attention_bias = True
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)

    # -------------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ):
        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, _ = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = self.q_proj(hidden_states).view(B, N_q, H, d).transpose(1, 2).contiguous()
        K_new = self.k_proj(kv_input).view(B, -1, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V_new = self.v_proj(kv_input).view(B, -1, self.num_kv_heads, d).transpose(1, 2).contiguous()

        position_ids = kwargs.get("position_ids")
        if position_ids is None and not is_cross_attention:
            past_len = 0
            if past_key_value is not None and past_key_value is not False:
                if hasattr(past_key_value, "get_seq_length"):
                    try:
                        past_len = past_key_value.get_seq_length(self.layer_idx)
                    except Exception:
                        try:
                            past_len = past_key_value.get_seq_length()
                        except Exception:
                            past_len = 0
                elif (
                    isinstance(past_key_value, (tuple, list))
                    and len(past_key_value) >= 2
                    and past_key_value[0] is not None
                ):
                    past_len = past_key_value[0].shape[-2]
            position_ids = (
                torch.arange(past_len, past_len + N_q, device=device, dtype=torch.long)
                .unsqueeze(0)
                .expand(B, -1)
            )

        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
            Q, K_new = apply_rotary_pos_emb(Q, K_new, cos, sin, position_ids=position_ids)
        elif self.rotary_emb is not None:
            if position_ids is None:
                position_ids = (
                    torch.arange(N_q, device=device, dtype=torch.long).unsqueeze(0).expand(B, -1)
                )
            try:
                cos, sin = self.rotary_emb(K_new, position_ids=position_ids)
            except TypeError:
                try:
                    cos, sin = self.rotary_emb(K_new, position_ids)
                except Exception:
                    cos, sin = None, None
            except Exception:
                cos, sin = None, None
            if cos is not None and sin is not None:
                Q, K_new = apply_rotary_pos_emb(Q, K_new, cos, sin, position_ids=position_ids)

        new_past_key_value: Any | None = None
        use_cache = kwargs.get("use_cache", False) or past_key_value is not None

        if is_cross_attention:
            K, V = K_new, V_new
            new_past_key_value = past_key_value
        else:
            if past_key_value is not None and past_key_value is not False:
                if hasattr(past_key_value, "update"):
                    try:
                        K, V, new_past_key_value = past_key_value.update(
                            K_new, V_new, self.layer_idx
                        )
                    except Exception:
                        K = torch.cat([past_key_value[0], K_new], dim=-2)
                        V = torch.cat([past_key_value[1], V_new], dim=-2)
                        new_past_key_value = (K, V)
                elif isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
                    past_k, past_v = past_key_value[0], past_key_value[1]
                    if past_k is not None:
                        K = torch.cat([past_k, K_new], dim=-2)
                        V = torch.cat([past_v, V_new], dim=-2)
                    else:
                        K, V = K_new, V_new
                    new_past_key_value = (K, V)
                else:
                    K, V = K_new, V_new
                    new_past_key_value = (K, V)
            else:
                K, V = K_new, V_new
                if use_cache:
                    new_past_key_value = (K, V)

            if self.rolling_cache_size is not None and K.shape[-2] > self.rolling_cache_size:
                K = K[..., -self.rolling_cache_size :, :].contiguous()
                V = V[..., -self.rolling_cache_size :, :].contiguous()
                if isinstance(new_past_key_value, tuple):
                    new_past_key_value = (K, V)

        N_k = K.shape[-2]

        if self.kv_repeat > 1:
            if K.shape[1] == self.num_kv_heads:
                K_full = repeat_kv(K, self.kv_repeat)
                V_full = repeat_kv(V, self.kv_repeat)
            elif K.shape[1] == self.num_heads:
                K_full, V_full = K, V
            else:
                raise ValueError(f"Unexpected K head dim {K.shape[1]}")
        else:
            K_full, V_full = K, V

        # ---------------- Scores ----------------
        scale = d**-0.5
        scores = torch.matmul(Q, K_full.transpose(-2, -1).contiguous()) * scale

        position_bias = kwargs.get("position_bias")
        if position_bias is None and self.has_relative_attention_bias:
            position_bias = self.compute_bias(N_q, N_k, device=device)
        if position_bias is not None:
            scores = scores + position_bias.to(scores.dtype)

        # =========================================================================
        # FIX: Apply Adaptive Span penalty BEFORE softmax to preserve logic & gradient
        # =========================================================================
        if not is_cross_attention:
            span = torch.sigmoid(self.span_logits) * self.max_span  # [H]
            if position_ids is not None:
                q_pos = position_ids[0].to(torch.float32)
                if N_k == N_q:
                    k_pos = q_pos
                else:
                    last_q = q_pos[-1].item()
                    k_pos = torch.arange(
                        last_q - N_k + 1, last_q + 1, device=device, dtype=torch.float32
                    )
            else:
                q_pos = torch.arange(N_q, device=device, dtype=torch.float32)
                k_pos = torch.arange(N_k, device=device, dtype=torch.float32)

            if getattr(self, "is_decoder", True) or kwargs.get("is_causal", False):
                dist = (q_pos.unsqueeze(1) - k_pos.unsqueeze(0)).clamp(min=0.0)
            else:
                dist = (q_pos.unsqueeze(1) - k_pos.unsqueeze(0)).abs()

            span_h = span.view(H, 1, 1).to(torch.float32)
            dist_mat = dist.unsqueeze(0)

            ramp = torch.clamp((self.ramp_size + span_h - dist_mat) / self.ramp_size, 0.0, 1.0)
            span_mask = torch.log(ramp.clamp(min=1e-9)).to(scores.dtype)
            scores = scores + span_mask.unsqueeze(0)

        # ---------------- Combined boolean mask ----------------
        is_causal = self.is_decoder and not is_cross_attention
        keep_mask = _build_keep_mask(
            B=B,
            N_q=N_q,
            N_k=N_k,
            is_causal=is_causal,
            attention_mask=attention_mask,
            encoder_attention_mask=(encoder_attention_mask if is_cross_attention else None),
            device=device,
        )

        fully_masked: torch.Tensor | None = None
        if keep_mask is not None:
            fully_masked = ~keep_mask.any(dim=-1, keepdim=True)
            mask_val = _safe_mask_value(scores.dtype)
            scores = scores.masked_fill(~keep_mask, mask_val)

        if fully_masked is not None:
            scores = torch.where(fully_masked, torch.zeros_like(scores), scores)

        if dtype in (torch.float16, torch.bfloat16) and scores.numel() > 2**30:
            attn = torch.empty_like(scores)
            chunk = max(64, min(512, N_q))
            for s in range(0, N_q, chunk):
                e = min(s + chunk, N_q)
                attn[:, :, s:e, :] = F.softmax(scores[:, :, s:e, :].float(), dim=-1).to(dtype)
        else:
            attn = F.softmax(scores.float(), dim=-1).to(dtype)

        if fully_masked is not None:
            attn = attn.masked_fill(fully_masked, 0.0)

        # ---------------- Dropout + head mask ----------------
        attn = self.dropout(attn)
        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn = attn * head_mask

        # ---------------- Context ----------------
        context = torch.matmul(attn, V_full)
        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)

        return self._format_output(
            output,
            attn if output_attentions else None,
            new_past_key_value,
            position_bias,
            output_attentions,
        )

    def _format_output(
        self,
        output: torch.Tensor,
        attn: torch.Tensor | None,
        past_kv: Any | None,
        position_bias: torch.Tensor | None,
        output_attentions: bool,
    ):
        style = self.return_style
        if style == "bert":
            return output, attn
        if style == "gpt2":
            if output_attentions:
                return output, past_kv, attn
            return output, past_kv
        if style == "t5":
            if output_attentions:
                return output, position_bias, attn
            return output, position_bias
        return output, attn, past_kv

Every 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_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 current_span : torch.Tensor

Expand source code

@property
def current_span(self) -> torch.Tensor:
    return torch.sigmoid(self.span_logits) * self.max_span

Methods

def compute_bias(self, query_length: int, key_length: int, device=None) ‑> torch.Tensor

Expand source code

def compute_bias(self, query_length: int, key_length: int, device=None) -> torch.Tensor:
    if self.relative_attention_bias is None:
        return torch.zeros((1, self.num_heads, query_length, key_length), device=device)
    ctx = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
    mem = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
    rel = mem - ctx
    buckets = self._relative_position_bucket(
        rel,
        bidirectional=(not self.is_decoder),
        num_buckets=self.relative_attention_num_buckets,
        max_distance=self.relative_attention_max_distance,
    )
    values = self.relative_attention_bias(buckets)
    return values.permute([2, 0, 1]).unsqueeze(0)

def get_span_regularization_loss(self, lambda_span: float = 0.01) ‑> torch.Tensor

Expand source code

def get_span_regularization_loss(self, lambda_span: float = 0.01) -> torch.Tensor:
    return lambda_span * torch.sigmoid(self.span_logits).mean()

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AdaptiveSparseAttention (config: Any | None = None,
rank: int | None = None,
sparsity_ratio: float | None = None,
dropout: float = 0.1,
lora_rank: int = 16,
comp_config: Any | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "adaptive_sparse",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "adaptive_sparse_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "adaptive_sparse_attn",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "adaptive_sparse_attention_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class AdaptiveSparseAttention(BaseAttention):
    is_custom_attention: bool = True
    supports_causal = True
    supports_bidirectional = True
    supports_gqa = True
    supports_rope = True
    triton_kernel: str = "sparse_attn"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        sparsity_ratio: float | None = None,
        dropout: float = 0.1,
        lora_rank: int = 16,
        comp_config: Any | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        # Resolve hidden_size and num_heads safely even if config is None
        hidden_size = kwargs.get("d_model") or kwargs.get("hidden_size") or kwargs.get("dim")
        num_heads = (
            kwargs.get("num_heads") or kwargs.get("n_heads") or kwargs.get("num_attention_heads")
        )
        num_kv_heads = kwargs.get("num_kv_heads") or kwargs.get("num_key_value_heads")

        head_dim = 0
        if config is not None:
            adapter = ModelAdapter(config)
            hidden_size = hidden_size or adapter.hidden_size
            num_heads = num_heads or adapter.num_heads
            num_kv_heads = num_kv_heads or adapter.num_key_value_heads
            head_dim = getattr(adapter, "head_dim", 0)

        if not hidden_size and source_attn is not None:
            cand = getattr(
                source_attn,
                "hidden_size",
                getattr(
                    source_attn,
                    "embed_dim",
                    getattr(source_attn, "all_head_size", None),
                ),
            )
            if (
                cand is None
                and hasattr(source_attn, "q_proj")
                and hasattr(source_attn.q_proj, "in_features")
            ):
                cand = source_attn.q_proj.in_features
            hidden_size = cand

        hidden_size = int(hidden_size or 768)
        num_heads = int(num_heads or 12)
        num_kv_heads = int(num_kv_heads or num_heads)
        head_dim = head_dim if head_dim > 0 else (hidden_size // num_heads if num_heads > 0 else 64)

        self.hidden_size: int = hidden_size
        self.d_model: int = hidden_size
        self.num_heads: int = num_heads
        self.num_kv_heads: int = num_kv_heads
        self.head_dim: int = 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 = max(1, self.num_heads // max(1, self.num_kv_heads))

        # Decoder & Causal detection
        model_type = getattr(config, "model_type", "").lower() if config else ""
        is_dec = (
            getattr(
                config,
                "is_decoder",
                getattr(config, "is_causal", kwargs.get("is_causal", False)),
            )
            if config
            else False
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(_mt.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)

        # Detect nested encoder architecture early
        source_cls = source_attn.__class__.__name__ if source_attn is not None else ""
        is_nested = source_cls in ("BertSelfAttention", "RobertaSelfAttention") or (
            source_cls.endswith(("SelfAttention", "SelfAttn"))
            and "distil" not in source_cls.lower()
        )
        has_out = (
            any(
                hasattr(source_attn, attr)
                for attr in (
                    "dense",
                    "out_proj",
                    "c_proj",
                    "out_lin",
                    "o_proj",
                    "wo",
                    "o",
                    "proj",
                )
            )
            if source_attn is not None
            else False
        )
        self.is_nested_encoder: bool = bool(
            kwargs.get("is_nested_encoder", (is_nested and not has_out))
        )

        # Head dim fallback from source_attn if head_dim <= 0 or not resolved properly
        if source_attn is not None and getattr(source_attn, "head_dim", 0):
            head_dim = source_attn.head_dim
            self.head_dim = head_dim
            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))

        # Default sparsity: lower for bidirectional encoders (0.1) than decoders (0.5)
        # Check explicit parameter, kwargs, comp_config, config
        cfg_encoder_sparsity = getattr(
            comp_config,
            "adaptive_sparse_encoder_sparsity",
            kwargs.get("adaptive_sparse_encoder_sparsity"),
        )
        if cfg_encoder_sparsity is None and config is not None:
            cfg_encoder_sparsity = getattr(config, "adaptive_sparse_encoder_sparsity", None)

        if sparsity_ratio is not None:
            effective_sparsity = float(sparsity_ratio)
        elif "sparsity_ratio" in kwargs:
            effective_sparsity = float(kwargs["sparsity_ratio"])
        elif comp_config is not None and hasattr(comp_config, "adaptive_sparsity_ratio"):
            effective_sparsity = float(comp_config.adaptive_sparsity_ratio)
        elif config is not None and hasattr(config, "adaptive_sparsity_ratio"):
            effective_sparsity = float(config.adaptive_sparsity_ratio)
        elif not self.is_decoder:
            if cfg_encoder_sparsity is not None:
                effective_sparsity = float(cfg_encoder_sparsity)
            else:
                effective_sparsity = 0.1
        else:
            effective_sparsity = 0.5

        self.default_sparsity: float = effective_sparsity

        # LoRA rank and alpha: independent reads
        if config is not None:
            self.lora_rank = getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", lora_rank)),
            )
            self.lora_alpha = getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        elif comp_config is not None:
            self.lora_rank = getattr(comp_config, "lora_rank", kwargs.get("lora_rank", lora_rank))
            self.lora_alpha = getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
        else:
            self.lora_rank = kwargs.get("lora_rank", lora_rank)
            self.lora_alpha = kwargs.get("lora_alpha", 32.0)
        self.freeze_base: bool = self.lora_rank > 0

        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=self.freeze_base,
        )
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )
        out_lora_rank = 0 if self.is_nested_encoder else self.lora_rank
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=out_lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Aliases
        self.q_lin = self.q_proj
        self.k_lin = self.k_proj
        self.v_lin = self.v_proj
        self.out_lin = self.out_proj

        # Per-head learnable sparsity logit initialized to target sparsity
        init_logit = math.log(
            max(1e-4, min(0.9999, effective_sparsity))
            / (1.0 - max(1e-4, min(0.9999, effective_sparsity)))
        )
        self.sparsity_logits = nn.Parameter(
            torch.full((self.num_heads,), fill_value=init_logit), requires_grad=True
        )
        self.output_scale = nn.Parameter(torch.ones(1))

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", dropout),
                ),
            )
            if config
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)
        self.source_attn = None
        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )

        self.has_relative_attention_bias: bool = False
        self.relative_attention_bias: nn.Embedding | None = None
        if config is not None:
            self.relative_attention_num_buckets: int = getattr(
                config, "relative_attention_num_buckets", 32
            )
            self.relative_attention_max_distance: int = getattr(
                config, "relative_attention_max_distance", 128
            )
        else:
            self.relative_attention_num_buckets = 32
            self.relative_attention_max_distance = 128

        self.triton_kwargs = {"block_size": 64, "eps": 1e-6}

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

    def freeze_base_weights(self) -> None:
        """Freeze base projection weights when LoRA adapters are present."""
        for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
            if hasattr(mod, "freeze_base_weights"):
                mod.freeze_base_weights()

    def unfreeze_base_weights(self) -> None:
        """Unfreeze base projection weights."""
        for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
            if hasattr(mod, "unfreeze_base_weights"):
                mod.unfreeze_base_weights()

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

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

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

    def _relative_position_bucket(
        self,
        relative_position: torch.Tensor,
        bidirectional: bool = True,
        num_buckets: int = 32,
        max_distance: int = 128,
    ) -> torch.Tensor:
        relative_buckets = 0
        if bidirectional:
            num_buckets //= 2
            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
            relative_position = torch.abs(relative_position)
        else:
            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
        max_exact = num_buckets // 2
        is_small = relative_position < max_exact
        relative_position_if_large = max_exact + (
            torch.log(relative_position.float() / max_exact)
            / math.log(max_distance / max_exact)
            * (num_buckets - max_exact)
        ).to(torch.long)
        relative_position_if_large = torch.min(
            relative_position_if_large,
            torch.full_like(relative_position_if_large, num_buckets - 1),
        )
        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
        return relative_buckets

    def compute_bias(
        self, query_length: int, key_length: int, device: torch.device | None = None
    ) -> torch.Tensor:
        if self.relative_attention_bias is not None:
            context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
            memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
            relative_position = memory_position - context_position
            relative_position_bucket = self._relative_position_bucket(
                relative_position,
                bidirectional=(not self.is_decoder),
                num_buckets=self.relative_attention_num_buckets,
                max_distance=self.relative_attention_max_distance,
            )
            values = self.relative_attention_bias(relative_position_bucket)
            values = values.permute([2, 0, 1]).unsqueeze(0)
            return values
        return torch.zeros((1, self.num_heads, query_length, key_length), device=device)

    def copy_weights(self, src: Any) -> None:
        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. HuggingFace / standard 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 / c_proj (handling Conv1D and LoRALinear wrapper)
        if (
            hasattr(src, "c_attn")
            and src.c_attn is not None
            and (
                extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
            )
        ):
            c_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (weight.ndim == 2 and weight.shape[0] < weight.shape[1])
                ):
                    weight = weight.t().contiguous()
                split_size = weight.shape[0] // 3
                bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias_mod.data.clone() if bias_mod is not None else None
                extracted["q"] = (
                    weight[:split_size, :],
                    bias_data[:split_size] if bias_data is not None else None,
                )
                extracted["k"] = (
                    weight[split_size : 2 * split_size, :],
                    (bias_data[split_size : 2 * split_size] if bias_data is not None else None),
                )
                extracted["v"] = (
                    weight[2 * split_size :, :],
                    bias_data[2 * split_size :] if bias_data is not None else None,
                )

        if hasattr(src, "c_proj") and src.c_proj is not None and extracted["o"][0] is None:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                if (
                    (c_proj_mod.__class__.__name__ == "Conv1D")
                    or (base_o.__class__.__name__ == "Conv1D")
                    or (w_o.ndim == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                extracted["o"] = (
                    w_o,
                    bias_o.data.clone() if bias_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
            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") and extracted["o"][0] is None:
            dense = src.output.dense
            extracted["o"] = (
                dense.weight.data.clone(),
                dense.bias.data.clone() if dense.bias is not None else None,
            )

        # 4. Modern Decoders (Qwen3.5, LLaMA, Mistral) & DistilBERT
        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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.clone(),
                                getattr(layer, "bias", None),
                            )
                            break

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

        # Check required weights
        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:
            raise RuntimeError(
                f"AdaptiveSparseAttention missing required weights for projections: {missing}"
            )

        # Frobenius norm validation on slices
        for proj_key, target_mod in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
        ]:
            w = extracted[proj_key][0]
            if w is not None:
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))
                norm = float(torch.norm(w[:out_d, :in_d].float()).item())
                if norm < 1e-4:
                    raise RuntimeError(
                        f"AdaptiveSparseAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm:.6f})"
                    )

        if extracted["o"][0] is not None:
            w_o = extracted["o"][0]
            out_d = min(self.out_proj.weight.size(0), w_o.size(0))
            in_d = min(self.out_proj.weight.size(1), w_o.size(1))
            norm_o = float(torch.norm(w_o[:out_d, :in_d].float()).item())
            if norm_o < 1e-4:
                raise RuntimeError(
                    f"AdaptiveSparseAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o:.6f})"
                )

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

        if hasattr(src, "rotary_emb") and src.rotary_emb is not None:
            self.rotary_emb = src.rotary_emb

        if hasattr(src, "has_relative_attention_bias"):
            self.has_relative_attention_bias = bool(src.has_relative_attention_bias)
        if hasattr(src, "relative_attention_bias") and src.relative_attention_bias is not None:
            self.relative_attention_bias = src.relative_attention_bias
            self.has_relative_attention_bias = True
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)

        if self.lora_rank > 0 or self.freeze_base:
            self.freeze_base_weights()

    def get_sparsity_loss(self) -> torch.Tensor:
        r"""Regularization loss encouraging head sparsity."""
        return torch.sigmoid(self.sparsity_logits).mean()

    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,
        *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

        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 RoPE (accept from position_embeddings or fallback to rotary_emb)
        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin, kwargs.get("position_ids"))
        elif self.rotary_emb is not None:
            rotary_mod = self.rotary_emb
            pos_ids = kwargs.get("position_ids")
            if pos_ids is None:
                pos_ids = torch.arange(N_q, device=device).unsqueeze(0).expand(B, -1)
            # 1. Try rotary_module(Q, position_ids)
            # 2. Try rotary_module(Q, seq_len=...)
            # 3. Try rotary_module(K, ...)
            cos, sin = None, None
            for probe_t in (Q, K):
                try:
                    res = rotary_mod(probe_t, pos_ids)
                    if isinstance(res, (tuple, list)) and len(res) >= 2:
                        cos, sin = res[0], res[1]
                        break
                except Exception as e:
                    logger.warning("RoPE embedding failed on (probe_t, pos_ids): %s", e)
                try:
                    res = rotary_mod(probe_t, seq_len=N_k)
                    if isinstance(res, (tuple, list)) and len(res) >= 2:
                        cos, sin = res[0], res[1]
                        break
                except Exception as e:
                    logger.warning("RoPE embedding failed on (probe_t, seq_len): %s", e)
                try:
                    res = rotary_mod(probe_t)
                    if isinstance(res, (tuple, list)) and len(res) >= 2:
                        cos, sin = res[0], res[1]
                        break
                except Exception as e:
                    logger.warning("RoPE embedding failed on (probe_t): %s", e)
            if cos is not None and sin is not None:
                Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids=pos_ids)

        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        chunk_size = 64
        scale = d**-0.5

        # Handle T5 position bias
        position_bias = kwargs.get("position_bias")
        if position_bias is None and self.has_relative_attention_bias:
            position_bias = self.compute_bias(N_q, N_k, device=device)

        mask = attention_mask if attention_mask is not None else kwargs.get("mask")

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

        sparsity_per_head = torch.sigmoid(self.sparsity_logits).clamp(0.0, 0.95)  # [H]
        k_per_head = torch.clamp(((1.0 - sparsity_per_head) * N_k).long(), min=1, max=N_k)  # [H]
        k_indices_base = (k_per_head - 1).view(1, H, 1, 1)

        valid_m = None
        if mask is not None and mask.dim() != 4:
            valid_m = get_token_valid_mask(mask, B, N_k, device)
            if valid_m is None:
                valid_m = torch.ones(B, N_k, device=device, dtype=torch.bool)

        context_chunks = []
        attn_chunks = [] if output_attentions else None

        for q_start in range(0, N_q, chunk_size):
            q_end = min(q_start + chunk_size, N_q)
            q_chunk = Q[:, :, q_start:q_end, :]  # [B, H, C_q, d]
            C_q = q_end - q_start

            scores_chunk = (
                torch.matmul(q_chunk, K.transpose(-2, -1).contiguous()) * scale
            )  # [B, H, C_q, N_k]

            if position_bias is not None:
                pb_chunk = position_bias[:, :, q_start:q_end, :]
                if mask is not None and mask.dim() == 4:
                    pb_chunk = pb_chunk + mask[:, :, q_start:q_end, :].to(scores_chunk.dtype)
                scores_chunk = scores_chunk + pb_chunk.to(scores_chunk.dtype)
            elif mask is not None:
                if mask.dim() == 4:
                    scores_chunk = scores_chunk + mask[:, :, q_start:q_end, :].to(
                        scores_chunk.dtype
                    )
                elif valid_m is not None:
                    scores_chunk = scores_chunk.masked_fill(
                        (~valid_m).unsqueeze(1).unsqueeze(2), -10000.0
                    )

            if is_decoder:
                q_idx = torch.arange(q_start, q_end, device=device).view(1, 1, C_q, 1)
                k_idx = torch.arange(N_k, device=device).view(1, 1, 1, N_k)
                causal_chunk = k_idx <= q_idx
                scores_chunk = scores_chunk.masked_fill(~causal_chunk, -10000.0)

            # Vectorized Top-K Thresholding on chunk with differentiable soft sigmoid threshold
            if not is_cross_attention:
                sorted_scores, _ = torch.sort(scores_chunk, dim=-1, descending=True)
                k_indices = k_indices_base.expand(B, H, C_q, 1)
                topk_thresh = torch.gather(sorted_scores, dim=-1, index=k_indices)
                temperature = 0.1
                soft_mask = torch.sigmoid((scores_chunk - topk_thresh) / temperature)
                scores_chunk = scores_chunk + torch.log(soft_mask + 1e-8)

            attn_chunk = torch.nan_to_num(F.softmax(scores_chunk.float(), dim=-1), nan=0.0).to(
                dtype
            )
            attn_chunk = self.dropout(attn_chunk)

            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.unsqueeze(1)
                elif head_mask.dim() == 4:
                    hm = head_mask[:, :, q_start:q_end, :]
                else:
                    hm = head_mask
                attn_chunk = attn_chunk * hm

            ctx_chunk = torch.matmul(attn_chunk, V)  # [B, H, C_q, d]
            context_chunks.append(ctx_chunk)
            if output_attentions:
                attn_chunks.append(attn_chunk)

        context = torch.cat(context_chunks, dim=2)
        context = context.transpose(1, 2).contiguous().view(B, N_q, self.all_head_size)
        attn = torch.cat(attn_chunks, dim=2) if output_attentions else None

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

        if position_bias is not None:
            return output * self.output_scale, position_bias
        return output * self.output_scale, (attn if output_attentions else None)

Every 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_bidirectional

The type of the None singleton.

var supports_causal

The type of the None singleton.

var supports_gqa

The type of the None singleton.

var supports_rope

The type of the None singleton.

var triton_kernel : str

The type of the None singleton.

Instance variables

prop key : nn.Module

Expand source code

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

prop query : nn.Module

Expand source code

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

prop value : nn.Module

Expand source code

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

Methods

def compute_bias(self, query_length: int, key_length: int, device: torch.device | None = None) ‑> torch.Tensor

Expand source code

def compute_bias(
    self, query_length: int, key_length: int, device: torch.device | None = None
) -> torch.Tensor:
    if self.relative_attention_bias is not None:
        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
        relative_position = memory_position - context_position
        relative_position_bucket = self._relative_position_bucket(
            relative_position,
            bidirectional=(not self.is_decoder),
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )
        values = self.relative_attention_bias(relative_position_bucket)
        values = values.permute([2, 0, 1]).unsqueeze(0)
        return values
    return torch.zeros((1, self.num_heads, query_length, key_length), device=device)

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    """Freeze base projection weights when LoRA adapters are present."""
    for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
        if hasattr(mod, "freeze_base_weights"):
            mod.freeze_base_weights()

Freeze base projection weights when LoRA adapters are present.

def get_sparsity_loss(self) ‑> torch.Tensor

Expand source code

def get_sparsity_loss(self) -> torch.Tensor:
    r"""Regularization loss encouraging head sparsity."""
    return torch.sigmoid(self.sparsity_logits).mean()

Regularization loss encouraging head sparsity.

def unfreeze_base_weights(self) ‑> None

Expand source code

def unfreeze_base_weights(self) -> None:
    """Unfreeze base projection weights."""
    for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
        if hasattr(mod, "unfreeze_base_weights"):
            mod.unfreeze_base_weights()

Unfreeze base projection weights.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AegisAttention (model_config: Any = None,
config: Any = None,
aegis_cfg: AegisConfigV5 | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("aegis_v5_module", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("AegisAttention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v5", supports_causal=True, supports_gqa=True)
class AegisAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        model_config: Any = None,
        config: Any = None,
        aegis_cfg: AegisConfigV5 | None = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = model_config if model_config is not None else config
        super().__init__(cfg, **kwargs)

        self.hidden_size: int = kwargs.get(
            "hidden_size",
            kwargs.get(
                "d_model",
                (
                    getattr(
                        cfg,
                        "hidden_size",
                        getattr(cfg, "d_model", getattr(cfg, "dim", 768)),
                    )
                    if cfg is not None
                    else 768
                ),
            ),
        )
        self.num_heads: int = kwargs.get(
            "num_heads",
            kwargs.get(
                "num_attention_heads",
                (
                    getattr(cfg, "num_attention_heads", getattr(cfg, "num_heads", 4))
                    if cfg is not None
                    else 4
                ),
            ),
        )
        self.head_dim: int = kwargs.get(
            "head_dim",
            (
                getattr(
                    cfg,
                    "head_dim",
                    getattr(cfg, "d_kv", self.hidden_size // max(1, self.num_heads)),
                )
                if cfg is not None
                else self.hidden_size // max(1, self.num_heads)
            ),
        )
        self.num_kv_heads: int = kwargs.get(
            "num_kv_heads",
            kwargs.get(
                "num_key_value_heads",
                (
                    getattr(
                        cfg,
                        "num_key_value_heads",
                        getattr(cfg, "num_kv_heads", self.num_heads),
                    )
                    if cfg is not None
                    else self.num_heads
                ),
            ),
        )
        if self.num_heads % max(1, self.num_kv_heads) != 0:
            raise ValueError(
                f"num_heads ({self.num_heads}) must be divisible by "
                f"num_kv_heads ({self.num_kv_heads})."
            )
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.rank: int = int(kwargs.pop("rank", 192))
        self.all_head_size: int = self.num_heads * self.head_dim
        self.all_kv_head_size: int = self.num_kv_heads * self.head_dim
        self.cfg: AegisConfigV5 = aegis_cfg if aegis_cfg is not None else AegisConfigV5()

        self.is_decoder: bool = _detect_decoder(cfg, is_decoder, pretrained_attn)
        self.is_nested_encoder: bool = False

        # Split-head allocation
        if self.cfg.num_local_heads + self.cfg.num_global_heads != self.num_heads:
            self.num_local_heads = (self.num_heads * 2) // 3
            self.num_global_heads = self.num_heads - self.num_local_heads
        else:
            self.num_local_heads = self.cfg.num_local_heads
            self.num_global_heads = self.cfg.num_global_heads
        if self.num_local_heads <= 0:
            self.num_local_heads = max(1, self.num_heads // 2)
            self.num_global_heads = self.num_heads - self.num_local_heads

        # Projections
        self.q_proj = FactorisedLinear(
            self.hidden_size,
            self.all_head_size,
            energy=self.cfg.svd_energy,
            min_rank=self.cfg.min_rank,
            max_rank=self.cfg.max_rank,
            rank=self.rank,
        )
        self.k_proj = FactorisedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            energy=self.cfg.svd_energy,
            min_rank=self.cfg.min_rank,
            max_rank=self.cfg.max_rank,
            rank=self.rank,
        )
        self.v_proj = FactorisedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            energy=self.cfg.svd_energy,
            min_rank=self.cfg.min_rank,
            max_rank=self.cfg.max_rank,
            rank=self.rank,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            self.hidden_size,
            energy=self.cfg.svd_energy,
            min_rank=self.cfg.min_rank,
            max_rank=self.cfg.max_rank,
            rank=self.rank,
        )

        # RFF projection for global heads (V5-FIX-14: buffer by default)
        rff_init = torch.empty(self.num_global_heads, self.head_dim, self.cfg.rff_features)
        for h in range(self.num_global_heads):
            dim = max(self.head_dim, self.cfg.rff_features)
            q_mat, _ = torch.linalg.qr(torch.randn(dim, dim))
            rff_init[h] = q_mat[: self.head_dim, : self.cfg.rff_features] / math.sqrt(
                self.cfg.rff_features
            )
        if self.cfg.rff_trainable:
            self.rff_proj = nn.Parameter(rff_init)
        else:
            self.register_buffer("rff_proj", rff_init)
        self.rff_features = self.cfg.rff_features

        # Adaptive gate (V5-FIX-10/11: only local heads need the blend).
        if self.cfg.use_adaptive_gate:
            self.blend_logits = nn.Parameter(torch.zeros(self.num_local_heads))
            self.gate_bias: float = float(self.cfg.gate_init_bias)
        else:
            self.register_parameter("blend_logits", None)
            self.gate_bias = 0.0

        self.dropout = nn.Dropout(self.cfg.dropout)

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

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

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

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

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

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

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

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

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

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

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

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

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

        # GPT-2 fused c_attn (Conv1D)
        if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            weight = old_attn.c_attn.weight.data
            bias = getattr(old_attn.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                weight = weight.t().contiguous()
            split = weight.shape[0] // 3
            extracted["q"] = (
                weight[:split, :],
                bias_data[:split] if bias_data is not None else None,
            )
            extracted["k"] = (
                weight[split : 2 * split, :],
                bias_data[split : 2 * split] if bias_data is not None else None,
            )
            extracted["v"] = (
                weight[2 * split :, :],
                bias_data[2 * split :] if bias_data is not None else None,
            )

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

        # BERT / RoBERTa nested encoder
        if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

        # Modern decoders / DistilBERT named mapping
        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        # Nested encoder detection
        cls_name = old_attn.__class__.__name__
        has_out_proj = any(
            hasattr(old_attn, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        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:
            raise RuntimeError(
                f"AegisAttention.copy_weights: could not extract mandatory projections "
                f"{missing} from {old_attn.__class__.__name__}."
            )

        for proj_key in ("q", "k", "v"):
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float()).item()
                if norm < 1e-4:
                    logger.warning(
                        "AegisAttention: extracted weight for '%s' has near-zero "
                        "Frobenius norm (%.2e).",
                        proj_key,
                        norm,
                    )
        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float()).item()
            if norm_o < 1e-4:
                logger.warning(
                    "AegisAttention: extracted weight for output projection has "
                    "near-zero Frobenius norm (%.2e).",
                    norm_o,
                )

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

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    target.decompose_from_weight(_orient(w, in_dim), b, self.cfg.svd_energy)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None, self.cfg.svd_energy)

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

    # ------------------------------------------------------------------
    # Mask handling (V5-FIX-18: centralized)
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """
        Return a boolean tensor (B, N_kv): True = key is valid (padding).
        Returns None when no mask is provided.
        """
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            # (B, 1, N_q, N_kv) -> OR over query dim -> (B, N_kv)
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        Hkv = self.num_kv_heads
        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]

        # --- Projections ---
        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
        V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

        # --- RoPE with past offset ---
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            past_k_probe = None
            if past_key_value is not None:
                if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1:
                    past_k_probe = past_key_value[0]
                elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                    past_k_probe = past_key_value.key_cache[0]
            past_len = 0 if past_k_probe is None else past_k_probe.shape[-2]
            cos = cos[..., past_len : past_len + N, :]
            sin = sin[..., past_len : past_len + N, :]
            Q, K_new = apply_rotary_pos_emb(Q, K_new, cos, sin)

        # --- KV cache concat ---
        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 = getattr(self, "layer_idx", 0)
                idx = min(idx, len(past_key_value.key_cache) - 1)
                past_k, past_v = (
                    past_key_value.key_cache[idx],
                    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

        # Update DynamicCache in place if applicable (V5-FIX-K2 semantic).
        if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
            idx = getattr(self, "layer_idx", 0)
            past_key_value.update(K_new, V_new, idx)
        else:
            pass

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

        # --- GQA expansion ---
        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        # --- Mask ---
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

        # --- Decide path: SDPA fast path vs split-head ---
        use_sdpa = (
            self.cfg.use_sdpa_fast_path
            and N_kv <= max(self.cfg.local_window, 32)
            and not output_attentions
        )

        if use_sdpa:
            # ------------------------------------------------------------------
            # SDPA fast path
            # ------------------------------------------------------------------
            q_sdp = Q.half() if Q.dtype == torch.float32 else Q
            k_sdp = K_full.half() if K_full.dtype == torch.float32 else K_full
            v_sdp = V_full.half() if V_full.dtype == torch.float32 else V_full

            attn_bias_sdp = None
            if eff_mask is not None and eff_mask.dim() == 4:
                attn_bias_sdp = eff_mask.to(q_sdp.dtype)
            elif eff_mask is not None and eff_mask.dim() == 2:
                attn_bias_sdp = ((eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0).to(
                    q_sdp.dtype
                )
            is_causal_sdp = self.is_decoder and not is_cross and N > 1

            out = F.scaled_dot_product_attention(
                q_sdp,
                k_sdp,
                v_sdp,
                attn_mask=attn_bias_sdp,
                dropout_p=self.dropout.p if self.training else 0.0,
                is_causal=is_causal_sdp,
            )
            out = out.to(hidden_states.dtype)
            attn_weights_out = None
        else:
            # ------------------------------------------------------------------
            # Split-head path
            # ------------------------------------------------------------------
            Hl, _Hg = self.num_local_heads, self.num_global_heads
            Q_local, Q_global = Q[:, :Hl], Q[:, Hl:]
            K_local, K_global = K_full[:, :Hl], K_full[:, Hl:]
            V_local, V_global = V_full[:, :Hl], V_full[:, Hl:]

            # Local heads: sliding-window attention.
            ctx_local = local_window_attention(
                Q_local,
                K_local,
                V_local,
                window=self.cfg.local_window,
                softpick_temperature=self.cfg.softpick_temperature,
                token_valid=token_valid,
                is_decoder=self.is_decoder and not is_cross,
                past_len=past_len,
            )

            # Global heads: kernel linear attention.
            Q_g_norm = F.normalize(Q_global.float(), p=2, dim=-1)
            K_g_norm = F.normalize(K_global.float(), p=2, dim=-1)
            if self.cfg.use_zero_sum_keys:
                K_g_norm = K_g_norm - K_g_norm.mean(dim=2, keepdim=True)

            rff_proj_f = self.rff_proj.float()
            max_clamp = 11.0
            Q_proj = torch.einsum("bhnd,hdf->bhnf", Q_g_norm, rff_proj_f).clamp(max=max_clamp)
            K_proj = torch.einsum("bhnd,hdf->bhnf", K_g_norm, rff_proj_f).clamp(max=max_clamp)
            scale_rff = 1.0 / math.sqrt(self.rff_features)
            Q_feat = torch.exp(Q_proj) * scale_rff
            K_feat = torch.exp(K_proj) * scale_rff

            # Zero-out invalid keys.
            if token_valid is not None:
                kv_valid = token_valid[:, None, :, None].to(K_feat.dtype)
                K_feat = K_feat * kv_valid
                V_global_f = V_global.float() * kv_valid.to(V_global.dtype)
            else:
                V_global_f = V_global.float()

            if self.is_decoder and not is_cross:
                ctx_global = causal_kernel_linear_attention(
                    Q_feat,
                    K_feat,
                    V_global_f,
                    chunk_size=self.cfg.causal_chunk_size,
                ).to(Q.dtype)
            else:
                KV = torch.einsum("bhnf,bhnd->bhfd", K_feat, V_global_f)
                K_sum = K_feat.sum(dim=2)
                denom = torch.einsum("bhnf,bhf->bhn", Q_feat, K_sum).unsqueeze(-1).clamp(min=1e-5)
                ctx_global = torch.einsum("bhnf,bhfd->bhnd", Q_feat, KV) / denom
                ctx_global = torch.nan_to_num(ctx_global, nan=0.0).to(Q.dtype)

            # Adaptive gate (V5-FIX-10): blend local heads with global mean.
            if self.blend_logits is not None:
                global_mean = ctx_global.mean(dim=1, keepdim=True)  # (B, 1, N, D)
                gate_local = torch.sigmoid(
                    self.blend_logits.view(1, Hl, 1, 1) + self.gate_bias
                )  # (1, Hl, 1, 1)
                ctx_local = gate_local * ctx_local + (1.0 - gate_local) * global_mean.to(
                    ctx_local.dtype
                )
            context = torch.cat([ctx_local, ctx_global], dim=1)

            # Reshape to (B, N, H, D)
            out = context.transpose(1, 2).contiguous().view(B, N, H * D)
            attn_weights_out = None  # not available in split-head path

        # --- head_mask (V5-FIX-01/19) ---
        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)
            # Normalize to (B, H, 1, 1) whether `out` is (B, H, N, D) or (B, N, H, D).
            if out.dim() == 4 and out.shape[1] == H:
                out = out * head_mask.to(out.dtype)
            else:
                # out is (B, N, H*D) — reshape
                out = out.view(B, N, H, D) * head_mask.to(out.dtype)
                out = out.reshape(B, N, H * D)

        # --- Ensure out is (B, N, H*D) ---
        if out.dim() == 4:
            if out.shape[1] == H:
                out = out.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
            else:
                out = out.contiguous().view(B, N, self.all_head_size)

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

        # --- Output projection (skipped for nested encoders) ---
        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)
        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)

        return output.to(hidden_states.dtype), attn_weights_out

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.zeros((), device=self.out_proj.up.weight.device)

Every 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.

Instance variables

prop c_proj : FactorisedLinear

Expand source code

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

prop dense : FactorisedLinear

Expand source code

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

prop k_lin : FactorisedLinear

Expand source code

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

prop key : FactorisedLinear

Expand source code

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

prop o_proj : FactorisedLinear

Expand source code

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

prop out_lin : FactorisedLinear

Expand source code

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

prop q_lin : FactorisedLinear

Expand source code

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

prop query : FactorisedLinear

Expand source code

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

prop v_lin : FactorisedLinear

Expand source code

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

prop value : FactorisedLinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    from neuralops.core.adapters import UniversalAdapter

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

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

    # GPT-2 fused c_attn (Conv1D)
    if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        weight = old_attn.c_attn.weight.data
        bias = getattr(old_attn.c_attn, "bias", None)
        bias_data = bias.data if bias is not None else None
        if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
            weight = weight.t().contiguous()
        split = weight.shape[0] // 3
        extracted["q"] = (
            weight[:split, :],
            bias_data[:split] if bias_data is not None else None,
        )
        extracted["k"] = (
            weight[split : 2 * split, :],
            bias_data[split : 2 * split] if bias_data is not None else None,
        )
        extracted["v"] = (
            weight[2 * split :, :],
            bias_data[2 * split :] if bias_data is not None else None,
        )

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

    # BERT / RoBERTa nested encoder
    if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        s = old_attn.self
        if hasattr(s, "query"):
            extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
        if hasattr(s, "key"):
            extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
        if hasattr(s, "value"):
            extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

    # Modern decoders / DistilBERT named mapping
    named_map = {
        "q": ["q_proj", "q_lin", "query"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    for proj, names in named_map.items():
        if extracted[proj][0] is None:
            for name in names:
                if hasattr(old_attn, name):
                    layer = getattr(old_attn, name)
                    if hasattr(layer, "weight"):
                        extracted[proj] = (
                            layer.weight.data,
                            getattr(layer, "bias", None),
                        )
                        break

    # Nested encoder detection
    cls_name = old_attn.__class__.__name__
    has_out_proj = any(
        hasattr(old_attn, a)
        for a in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    if (
        hasattr(old_attn, "out_lin")
        or hasattr(old_attn, "c_proj")
        or hasattr(old_attn, "o_proj")
    ):
        self.is_nested_encoder = False
    elif (
        hasattr(old_attn, "self")
        and hasattr(old_attn.self, "query")
        and hasattr(old_attn, "output")
        and hasattr(old_attn.output, "dense")
    ) or (
        cls_name in ("BertSelfAttention", "RobertaSelfAttention")
        or (cls_name.endswith("SelfAttention") and not has_out_proj)
    ):
        self.is_nested_encoder = True
    else:
        self.is_nested_encoder = False

    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:
        raise RuntimeError(
            f"AegisAttention.copy_weights: could not extract mandatory projections "
            f"{missing} from {old_attn.__class__.__name__}."
        )

    for proj_key in ("q", "k", "v"):
        w = extracted[proj_key][0]
        if w is not None:
            norm = torch.linalg.norm(w.float()).item()
            if norm < 1e-4:
                logger.warning(
                    "AegisAttention: extracted weight for '%s' has near-zero "
                    "Frobenius norm (%.2e).",
                    proj_key,
                    norm,
                )
    if extracted["o"][0] is not None:
        norm_o = torch.linalg.norm(extracted["o"][0].float()).item()
        if norm_o < 1e-4:
            logger.warning(
                "AegisAttention: extracted weight for output projection has "
                "near-zero Frobenius norm (%.2e).",
                norm_o,
            )

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

    with torch.no_grad():
        for proj_key, target in (
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ):
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                target.decompose_from_weight(_orient(w, in_dim), b, self.cfg.svd_energy)
            elif proj_key == "o" and self.is_nested_encoder:
                eye_w = torch.eye(
                    self.all_head_size,
                    self.all_head_size,
                    device=self.out_proj.up.weight.device,
                    dtype=self.out_proj.up.weight.dtype,
                )
                target.decompose_from_weight(eye_w, None, self.cfg.svd_energy)

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

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    H, D = self.num_heads, self.head_dim
    Hkv = self.num_kv_heads
    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]

    # --- Projections ---
    Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
    K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
    V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

    # --- RoPE with past offset ---
    if position_embeddings is not None and not is_cross:
        cos, sin = position_embeddings
        past_k_probe = None
        if past_key_value is not None:
            if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1:
                past_k_probe = past_key_value[0]
            elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                past_k_probe = past_key_value.key_cache[0]
        past_len = 0 if past_k_probe is None else past_k_probe.shape[-2]
        cos = cos[..., past_len : past_len + N, :]
        sin = sin[..., past_len : past_len + N, :]
        Q, K_new = apply_rotary_pos_emb(Q, K_new, cos, sin)

    # --- KV cache concat ---
    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 = getattr(self, "layer_idx", 0)
            idx = min(idx, len(past_key_value.key_cache) - 1)
            past_k, past_v = (
                past_key_value.key_cache[idx],
                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

    # Update DynamicCache in place if applicable (V5-FIX-K2 semantic).
    if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
        idx = getattr(self, "layer_idx", 0)
        past_key_value.update(K_new, V_new, idx)
    else:
        pass

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

    # --- GQA expansion ---
    K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
    V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

    # --- Mask ---
    eff_mask = encoder_attention_mask if is_cross else attention_mask
    token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

    # --- Decide path: SDPA fast path vs split-head ---
    use_sdpa = (
        self.cfg.use_sdpa_fast_path
        and N_kv <= max(self.cfg.local_window, 32)
        and not output_attentions
    )

    if use_sdpa:
        # ------------------------------------------------------------------
        # SDPA fast path
        # ------------------------------------------------------------------
        q_sdp = Q.half() if Q.dtype == torch.float32 else Q
        k_sdp = K_full.half() if K_full.dtype == torch.float32 else K_full
        v_sdp = V_full.half() if V_full.dtype == torch.float32 else V_full

        attn_bias_sdp = None
        if eff_mask is not None and eff_mask.dim() == 4:
            attn_bias_sdp = eff_mask.to(q_sdp.dtype)
        elif eff_mask is not None and eff_mask.dim() == 2:
            attn_bias_sdp = ((eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0).to(
                q_sdp.dtype
            )
        is_causal_sdp = self.is_decoder and not is_cross and N > 1

        out = F.scaled_dot_product_attention(
            q_sdp,
            k_sdp,
            v_sdp,
            attn_mask=attn_bias_sdp,
            dropout_p=self.dropout.p if self.training else 0.0,
            is_causal=is_causal_sdp,
        )
        out = out.to(hidden_states.dtype)
        attn_weights_out = None
    else:
        # ------------------------------------------------------------------
        # Split-head path
        # ------------------------------------------------------------------
        Hl, _Hg = self.num_local_heads, self.num_global_heads
        Q_local, Q_global = Q[:, :Hl], Q[:, Hl:]
        K_local, K_global = K_full[:, :Hl], K_full[:, Hl:]
        V_local, V_global = V_full[:, :Hl], V_full[:, Hl:]

        # Local heads: sliding-window attention.
        ctx_local = local_window_attention(
            Q_local,
            K_local,
            V_local,
            window=self.cfg.local_window,
            softpick_temperature=self.cfg.softpick_temperature,
            token_valid=token_valid,
            is_decoder=self.is_decoder and not is_cross,
            past_len=past_len,
        )

        # Global heads: kernel linear attention.
        Q_g_norm = F.normalize(Q_global.float(), p=2, dim=-1)
        K_g_norm = F.normalize(K_global.float(), p=2, dim=-1)
        if self.cfg.use_zero_sum_keys:
            K_g_norm = K_g_norm - K_g_norm.mean(dim=2, keepdim=True)

        rff_proj_f = self.rff_proj.float()
        max_clamp = 11.0
        Q_proj = torch.einsum("bhnd,hdf->bhnf", Q_g_norm, rff_proj_f).clamp(max=max_clamp)
        K_proj = torch.einsum("bhnd,hdf->bhnf", K_g_norm, rff_proj_f).clamp(max=max_clamp)
        scale_rff = 1.0 / math.sqrt(self.rff_features)
        Q_feat = torch.exp(Q_proj) * scale_rff
        K_feat = torch.exp(K_proj) * scale_rff

        # Zero-out invalid keys.
        if token_valid is not None:
            kv_valid = token_valid[:, None, :, None].to(K_feat.dtype)
            K_feat = K_feat * kv_valid
            V_global_f = V_global.float() * kv_valid.to(V_global.dtype)
        else:
            V_global_f = V_global.float()

        if self.is_decoder and not is_cross:
            ctx_global = causal_kernel_linear_attention(
                Q_feat,
                K_feat,
                V_global_f,
                chunk_size=self.cfg.causal_chunk_size,
            ).to(Q.dtype)
        else:
            KV = torch.einsum("bhnf,bhnd->bhfd", K_feat, V_global_f)
            K_sum = K_feat.sum(dim=2)
            denom = torch.einsum("bhnf,bhf->bhn", Q_feat, K_sum).unsqueeze(-1).clamp(min=1e-5)
            ctx_global = torch.einsum("bhnf,bhfd->bhnd", Q_feat, KV) / denom
            ctx_global = torch.nan_to_num(ctx_global, nan=0.0).to(Q.dtype)

        # Adaptive gate (V5-FIX-10): blend local heads with global mean.
        if self.blend_logits is not None:
            global_mean = ctx_global.mean(dim=1, keepdim=True)  # (B, 1, N, D)
            gate_local = torch.sigmoid(
                self.blend_logits.view(1, Hl, 1, 1) + self.gate_bias
            )  # (1, Hl, 1, 1)
            ctx_local = gate_local * ctx_local + (1.0 - gate_local) * global_mean.to(
                ctx_local.dtype
            )
        context = torch.cat([ctx_local, ctx_global], dim=1)

        # Reshape to (B, N, H, D)
        out = context.transpose(1, 2).contiguous().view(B, N, H * D)
        attn_weights_out = None  # not available in split-head path

    # --- head_mask (V5-FIX-01/19) ---
    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)
        # Normalize to (B, H, 1, 1) whether `out` is (B, H, N, D) or (B, N, H, D).
        if out.dim() == 4 and out.shape[1] == H:
            out = out * head_mask.to(out.dtype)
        else:
            # out is (B, N, H*D) — reshape
            out = out.view(B, N, H, D) * head_mask.to(out.dtype)
            out = out.reshape(B, N, H * D)

    # --- Ensure out is (B, N, H*D) ---
    if out.dim() == 4:
        if out.shape[1] == H:
            out = out.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        else:
            out = out.contiguous().view(B, N, self.all_head_size)

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

    # --- Output projection (skipped for nested encoders) ---
    if self.is_nested_encoder:
        output = out
    else:
        output = self.out_proj(out)
    output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)

    return output.to(hidden_states.dtype), attn_weights_out

Must return (output, attention_weights_or_None).

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.zeros((), device=self.out_proj.up.weight.device)

class AegisAttentionV5 (model_config: Any = None,
config: Any = None,
aegis_cfg: AegisConfigV5 | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("aegis_v5_module", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("AegisAttention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v5", supports_causal=True, supports_gqa=True)
class AegisAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        model_config: Any = None,
        config: Any = None,
        aegis_cfg: AegisConfigV5 | None = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = model_config if model_config is not None else config
        super().__init__(cfg, **kwargs)

        self.hidden_size: int = kwargs.get(
            "hidden_size",
            kwargs.get(
                "d_model",
                (
                    getattr(
                        cfg,
                        "hidden_size",
                        getattr(cfg, "d_model", getattr(cfg, "dim", 768)),
                    )
                    if cfg is not None
                    else 768
                ),
            ),
        )
        self.num_heads: int = kwargs.get(
            "num_heads",
            kwargs.get(
                "num_attention_heads",
                (
                    getattr(cfg, "num_attention_heads", getattr(cfg, "num_heads", 4))
                    if cfg is not None
                    else 4
                ),
            ),
        )
        self.head_dim: int = kwargs.get(
            "head_dim",
            (
                getattr(
                    cfg,
                    "head_dim",
                    getattr(cfg, "d_kv", self.hidden_size // max(1, self.num_heads)),
                )
                if cfg is not None
                else self.hidden_size // max(1, self.num_heads)
            ),
        )
        self.num_kv_heads: int = kwargs.get(
            "num_kv_heads",
            kwargs.get(
                "num_key_value_heads",
                (
                    getattr(
                        cfg,
                        "num_key_value_heads",
                        getattr(cfg, "num_kv_heads", self.num_heads),
                    )
                    if cfg is not None
                    else self.num_heads
                ),
            ),
        )
        if self.num_heads % max(1, self.num_kv_heads) != 0:
            raise ValueError(
                f"num_heads ({self.num_heads}) must be divisible by "
                f"num_kv_heads ({self.num_kv_heads})."
            )
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.rank: int = int(kwargs.pop("rank", 192))
        self.all_head_size: int = self.num_heads * self.head_dim
        self.all_kv_head_size: int = self.num_kv_heads * self.head_dim
        self.cfg: AegisConfigV5 = aegis_cfg if aegis_cfg is not None else AegisConfigV5()

        self.is_decoder: bool = _detect_decoder(cfg, is_decoder, pretrained_attn)
        self.is_nested_encoder: bool = False

        # Split-head allocation
        if self.cfg.num_local_heads + self.cfg.num_global_heads != self.num_heads:
            self.num_local_heads = (self.num_heads * 2) // 3
            self.num_global_heads = self.num_heads - self.num_local_heads
        else:
            self.num_local_heads = self.cfg.num_local_heads
            self.num_global_heads = self.cfg.num_global_heads
        if self.num_local_heads <= 0:
            self.num_local_heads = max(1, self.num_heads // 2)
            self.num_global_heads = self.num_heads - self.num_local_heads

        # Projections
        self.q_proj = FactorisedLinear(
            self.hidden_size,
            self.all_head_size,
            energy=self.cfg.svd_energy,
            min_rank=self.cfg.min_rank,
            max_rank=self.cfg.max_rank,
            rank=self.rank,
        )
        self.k_proj = FactorisedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            energy=self.cfg.svd_energy,
            min_rank=self.cfg.min_rank,
            max_rank=self.cfg.max_rank,
            rank=self.rank,
        )
        self.v_proj = FactorisedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            energy=self.cfg.svd_energy,
            min_rank=self.cfg.min_rank,
            max_rank=self.cfg.max_rank,
            rank=self.rank,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            self.hidden_size,
            energy=self.cfg.svd_energy,
            min_rank=self.cfg.min_rank,
            max_rank=self.cfg.max_rank,
            rank=self.rank,
        )

        # RFF projection for global heads (V5-FIX-14: buffer by default)
        rff_init = torch.empty(self.num_global_heads, self.head_dim, self.cfg.rff_features)
        for h in range(self.num_global_heads):
            dim = max(self.head_dim, self.cfg.rff_features)
            q_mat, _ = torch.linalg.qr(torch.randn(dim, dim))
            rff_init[h] = q_mat[: self.head_dim, : self.cfg.rff_features] / math.sqrt(
                self.cfg.rff_features
            )
        if self.cfg.rff_trainable:
            self.rff_proj = nn.Parameter(rff_init)
        else:
            self.register_buffer("rff_proj", rff_init)
        self.rff_features = self.cfg.rff_features

        # Adaptive gate (V5-FIX-10/11: only local heads need the blend).
        if self.cfg.use_adaptive_gate:
            self.blend_logits = nn.Parameter(torch.zeros(self.num_local_heads))
            self.gate_bias: float = float(self.cfg.gate_init_bias)
        else:
            self.register_parameter("blend_logits", None)
            self.gate_bias = 0.0

        self.dropout = nn.Dropout(self.cfg.dropout)

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

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

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

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

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

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

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

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

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

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

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

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

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

        # GPT-2 fused c_attn (Conv1D)
        if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            weight = old_attn.c_attn.weight.data
            bias = getattr(old_attn.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                weight = weight.t().contiguous()
            split = weight.shape[0] // 3
            extracted["q"] = (
                weight[:split, :],
                bias_data[:split] if bias_data is not None else None,
            )
            extracted["k"] = (
                weight[split : 2 * split, :],
                bias_data[split : 2 * split] if bias_data is not None else None,
            )
            extracted["v"] = (
                weight[2 * split :, :],
                bias_data[2 * split :] if bias_data is not None else None,
            )

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

        # BERT / RoBERTa nested encoder
        if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

        # Modern decoders / DistilBERT named mapping
        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        # Nested encoder detection
        cls_name = old_attn.__class__.__name__
        has_out_proj = any(
            hasattr(old_attn, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        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:
            raise RuntimeError(
                f"AegisAttention.copy_weights: could not extract mandatory projections "
                f"{missing} from {old_attn.__class__.__name__}."
            )

        for proj_key in ("q", "k", "v"):
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float()).item()
                if norm < 1e-4:
                    logger.warning(
                        "AegisAttention: extracted weight for '%s' has near-zero "
                        "Frobenius norm (%.2e).",
                        proj_key,
                        norm,
                    )
        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float()).item()
            if norm_o < 1e-4:
                logger.warning(
                    "AegisAttention: extracted weight for output projection has "
                    "near-zero Frobenius norm (%.2e).",
                    norm_o,
                )

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

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    target.decompose_from_weight(_orient(w, in_dim), b, self.cfg.svd_energy)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None, self.cfg.svd_energy)

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

    # ------------------------------------------------------------------
    # Mask handling (V5-FIX-18: centralized)
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """
        Return a boolean tensor (B, N_kv): True = key is valid (padding).
        Returns None when no mask is provided.
        """
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            # (B, 1, N_q, N_kv) -> OR over query dim -> (B, N_kv)
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        Hkv = self.num_kv_heads
        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]

        # --- Projections ---
        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
        V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

        # --- RoPE with past offset ---
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            past_k_probe = None
            if past_key_value is not None:
                if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1:
                    past_k_probe = past_key_value[0]
                elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                    past_k_probe = past_key_value.key_cache[0]
            past_len = 0 if past_k_probe is None else past_k_probe.shape[-2]
            cos = cos[..., past_len : past_len + N, :]
            sin = sin[..., past_len : past_len + N, :]
            Q, K_new = apply_rotary_pos_emb(Q, K_new, cos, sin)

        # --- KV cache concat ---
        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 = getattr(self, "layer_idx", 0)
                idx = min(idx, len(past_key_value.key_cache) - 1)
                past_k, past_v = (
                    past_key_value.key_cache[idx],
                    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

        # Update DynamicCache in place if applicable (V5-FIX-K2 semantic).
        if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
            idx = getattr(self, "layer_idx", 0)
            past_key_value.update(K_new, V_new, idx)
        else:
            pass

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

        # --- GQA expansion ---
        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        # --- Mask ---
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

        # --- Decide path: SDPA fast path vs split-head ---
        use_sdpa = (
            self.cfg.use_sdpa_fast_path
            and N_kv <= max(self.cfg.local_window, 32)
            and not output_attentions
        )

        if use_sdpa:
            # ------------------------------------------------------------------
            # SDPA fast path
            # ------------------------------------------------------------------
            q_sdp = Q.half() if Q.dtype == torch.float32 else Q
            k_sdp = K_full.half() if K_full.dtype == torch.float32 else K_full
            v_sdp = V_full.half() if V_full.dtype == torch.float32 else V_full

            attn_bias_sdp = None
            if eff_mask is not None and eff_mask.dim() == 4:
                attn_bias_sdp = eff_mask.to(q_sdp.dtype)
            elif eff_mask is not None and eff_mask.dim() == 2:
                attn_bias_sdp = ((eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0).to(
                    q_sdp.dtype
                )
            is_causal_sdp = self.is_decoder and not is_cross and N > 1

            out = F.scaled_dot_product_attention(
                q_sdp,
                k_sdp,
                v_sdp,
                attn_mask=attn_bias_sdp,
                dropout_p=self.dropout.p if self.training else 0.0,
                is_causal=is_causal_sdp,
            )
            out = out.to(hidden_states.dtype)
            attn_weights_out = None
        else:
            # ------------------------------------------------------------------
            # Split-head path
            # ------------------------------------------------------------------
            Hl, _Hg = self.num_local_heads, self.num_global_heads
            Q_local, Q_global = Q[:, :Hl], Q[:, Hl:]
            K_local, K_global = K_full[:, :Hl], K_full[:, Hl:]
            V_local, V_global = V_full[:, :Hl], V_full[:, Hl:]

            # Local heads: sliding-window attention.
            ctx_local = local_window_attention(
                Q_local,
                K_local,
                V_local,
                window=self.cfg.local_window,
                softpick_temperature=self.cfg.softpick_temperature,
                token_valid=token_valid,
                is_decoder=self.is_decoder and not is_cross,
                past_len=past_len,
            )

            # Global heads: kernel linear attention.
            Q_g_norm = F.normalize(Q_global.float(), p=2, dim=-1)
            K_g_norm = F.normalize(K_global.float(), p=2, dim=-1)
            if self.cfg.use_zero_sum_keys:
                K_g_norm = K_g_norm - K_g_norm.mean(dim=2, keepdim=True)

            rff_proj_f = self.rff_proj.float()
            max_clamp = 11.0
            Q_proj = torch.einsum("bhnd,hdf->bhnf", Q_g_norm, rff_proj_f).clamp(max=max_clamp)
            K_proj = torch.einsum("bhnd,hdf->bhnf", K_g_norm, rff_proj_f).clamp(max=max_clamp)
            scale_rff = 1.0 / math.sqrt(self.rff_features)
            Q_feat = torch.exp(Q_proj) * scale_rff
            K_feat = torch.exp(K_proj) * scale_rff

            # Zero-out invalid keys.
            if token_valid is not None:
                kv_valid = token_valid[:, None, :, None].to(K_feat.dtype)
                K_feat = K_feat * kv_valid
                V_global_f = V_global.float() * kv_valid.to(V_global.dtype)
            else:
                V_global_f = V_global.float()

            if self.is_decoder and not is_cross:
                ctx_global = causal_kernel_linear_attention(
                    Q_feat,
                    K_feat,
                    V_global_f,
                    chunk_size=self.cfg.causal_chunk_size,
                ).to(Q.dtype)
            else:
                KV = torch.einsum("bhnf,bhnd->bhfd", K_feat, V_global_f)
                K_sum = K_feat.sum(dim=2)
                denom = torch.einsum("bhnf,bhf->bhn", Q_feat, K_sum).unsqueeze(-1).clamp(min=1e-5)
                ctx_global = torch.einsum("bhnf,bhfd->bhnd", Q_feat, KV) / denom
                ctx_global = torch.nan_to_num(ctx_global, nan=0.0).to(Q.dtype)

            # Adaptive gate (V5-FIX-10): blend local heads with global mean.
            if self.blend_logits is not None:
                global_mean = ctx_global.mean(dim=1, keepdim=True)  # (B, 1, N, D)
                gate_local = torch.sigmoid(
                    self.blend_logits.view(1, Hl, 1, 1) + self.gate_bias
                )  # (1, Hl, 1, 1)
                ctx_local = gate_local * ctx_local + (1.0 - gate_local) * global_mean.to(
                    ctx_local.dtype
                )
            context = torch.cat([ctx_local, ctx_global], dim=1)

            # Reshape to (B, N, H, D)
            out = context.transpose(1, 2).contiguous().view(B, N, H * D)
            attn_weights_out = None  # not available in split-head path

        # --- head_mask (V5-FIX-01/19) ---
        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)
            # Normalize to (B, H, 1, 1) whether `out` is (B, H, N, D) or (B, N, H, D).
            if out.dim() == 4 and out.shape[1] == H:
                out = out * head_mask.to(out.dtype)
            else:
                # out is (B, N, H*D) — reshape
                out = out.view(B, N, H, D) * head_mask.to(out.dtype)
                out = out.reshape(B, N, H * D)

        # --- Ensure out is (B, N, H*D) ---
        if out.dim() == 4:
            if out.shape[1] == H:
                out = out.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
            else:
                out = out.contiguous().view(B, N, self.all_head_size)

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

        # --- Output projection (skipped for nested encoders) ---
        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)
        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)

        return output.to(hidden_states.dtype), attn_weights_out

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.zeros((), device=self.out_proj.up.weight.device)

Every 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.

Instance variables

prop c_proj : FactorisedLinear

Expand source code

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

prop dense : FactorisedLinear

Expand source code

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

prop k_lin : FactorisedLinear

Expand source code

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

prop key : FactorisedLinear

Expand source code

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

prop o_proj : FactorisedLinear

Expand source code

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

prop out_lin : FactorisedLinear

Expand source code

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

prop q_lin : FactorisedLinear

Expand source code

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

prop query : FactorisedLinear

Expand source code

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

prop v_lin : FactorisedLinear

Expand source code

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

prop value : FactorisedLinear

Expand source code

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

Methods

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.zeros((), device=self.out_proj.up.weight.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AegisAttentionConfigV2 (hidden_size: int = 768,
num_attention_heads: int = 12,
num_heads: int = 12,
head_dim: int = 64,
num_key_value_heads: int = 12,
local_window: int = 64,
rank: int = 384,
kernel_feature_dim: int = 128,
dropout: float = 0.1,
use_adaptive_pruning: bool = True,
gate_init_bias: float = 1.8,
is_decoder: bool = False)

Expand source code

@dataclass
class AegisAttentionConfigV2:
    hidden_size: int = 768
    num_attention_heads: int = 12
    num_heads: int = 12
    head_dim: int = 64
    num_key_value_heads: int = 12
    local_window: int = 64
    rank: int = 384
    kernel_feature_dim: int = 128
    dropout: float = 0.1
    use_adaptive_pruning: bool = True
    gate_init_bias: float = 1.8
    is_decoder: bool = False

AegisAttentionConfigV2(hidden_size: int = 768, num_attention_heads: int = 12, num_heads: int = 12, head_dim: int = 64, num_key_value_heads: int = 12, local_window: int = 64, rank: int = 384, kernel_feature_dim: int = 128, dropout: float = 0.1, use_adaptive_pruning: bool = True, gate_init_bias: float = 1.8, is_decoder: bool = False)

Instance variables

var dropout : float

The type of the None singleton.

var gate_init_bias : float

The type of the None singleton.

var head_dim : int

The type of the None singleton.

var hidden_size : int

The type of the None singleton.

var is_decoder : bool

The type of the None singleton.

var kernel_feature_dim : int

The type of the None singleton.

var local_window : int

The type of the None singleton.

var num_attention_heads : int

The type of the None singleton.

var num_heads : int

The type of the None singleton.

var num_key_value_heads : int

The type of the None singleton.

var rank : int

The type of the None singleton.

var use_adaptive_pruning : bool

The type of the None singleton.

class AegisConfigV2 (hidden_size: int = 768,
num_attention_heads: int = 12,
num_heads: int = 12,
head_dim: int = 64,
num_key_value_heads: int = 12,
local_window: int = 64,
rank: int = 384,
kernel_feature_dim: int = 128,
dropout: float = 0.1,
use_adaptive_pruning: bool = True,
gate_init_bias: float = 1.8,
is_decoder: bool = False)

Expand source code

@dataclass
class AegisAttentionConfigV2:
    hidden_size: int = 768
    num_attention_heads: int = 12
    num_heads: int = 12
    head_dim: int = 64
    num_key_value_heads: int = 12
    local_window: int = 64
    rank: int = 384
    kernel_feature_dim: int = 128
    dropout: float = 0.1
    use_adaptive_pruning: bool = True
    gate_init_bias: float = 1.8
    is_decoder: bool = False

AegisAttentionConfigV2(hidden_size: int = 768, num_attention_heads: int = 12, num_heads: int = 12, head_dim: int = 64, num_key_value_heads: int = 12, local_window: int = 64, rank: int = 384, kernel_feature_dim: int = 128, dropout: float = 0.1, use_adaptive_pruning: bool = True, gate_init_bias: float = 1.8, is_decoder: bool = False)

Instance variables

var dropout : float

The type of the None singleton.

var gate_init_bias : float

The type of the None singleton.

var head_dim : int

The type of the None singleton.

var hidden_size : int

The type of the None singleton.

var is_decoder : bool

The type of the None singleton.

var kernel_feature_dim : int

The type of the None singleton.

var local_window : int

The type of the None singleton.

var num_attention_heads : int

The type of the None singleton.

var num_heads : int

The type of the None singleton.

var num_key_value_heads : int

The type of the None singleton.

var rank : int

The type of the None singleton.

var use_adaptive_pruning : bool

The type of the None singleton.

class AegisAttentionV3 (config: Any = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("aegis_v2", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v1", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v3", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v4", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v1_module", supports_causal=True, supports_gqa=True)
class AegisAttentionV2(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        if config is None or not isinstance(config, AegisConfigV2):
            h_size = (
                kwargs.get(
                    "hidden_size",
                    getattr(
                        config,
                        "hidden_size",
                        getattr(config, "d_model", getattr(config, "dim", 768)),
                    ),
                )
                if config is not None
                else 768
            )
            n_heads = (
                kwargs.get(
                    "num_heads",
                    getattr(config, "num_attention_heads", getattr(config, "num_heads", 12)),
                )
                if config is not None
                else 12
            )
            h_dim = (
                kwargs.get(
                    "head_dim",
                    getattr(
                        config,
                        "head_dim",
                        getattr(config, "d_kv", h_size // max(1, n_heads)),
                    ),
                )
                if config is not None
                else h_size // max(1, n_heads)
            )
            n_kv_heads = (
                kwargs.get(
                    "num_key_value_heads",
                    getattr(
                        config,
                        "num_key_value_heads",
                        getattr(config, "num_kv_heads", n_heads),
                    ),
                )
                if config is not None
                else n_heads
            )

            # V2-FIX-01: robust decoder detection
            dec = _detect_decoder(config, is_decoder, pretrained_attn)

            config = AegisConfigV2(
                hidden_size=h_size,
                num_attention_heads=n_heads,
                num_heads=n_heads,
                num_key_value_heads=n_kv_heads,
                head_dim=h_dim,
                local_window=(
                    kwargs.get(
                        "local_window",
                        getattr(config, "block_size", getattr(config, "local_window", 64)),
                    )
                    if config is not None
                    else 64
                ),
                rank=(
                    kwargs.pop("rank", getattr(config, "lora_r", getattr(config, "rank", 64)))
                    if config is not None
                    else 64
                ),
                kernel_feature_dim=(
                    getattr(
                        config,
                        "rff_features",
                        getattr(config, "kernel_feature_dim", 128),
                    )
                    if config is not None
                    else 128
                ),
                dropout=(
                    kwargs.get("dropout", getattr(config, "dropout", 0.1))
                    if config is not None
                    else 0.1
                ),
                use_adaptive_pruning=(
                    getattr(config, "use_gate", True) if config is not None else True
                ),
                gate_init_bias=(
                    getattr(config, "gate_init_bias", 1.8) if config is not None else 1.8
                ),
                is_decoder=dec,
            )

        super().__init__(config)
        self.config = config
        self.hidden_size: int = config.hidden_size
        self.num_heads: int = config.num_attention_heads
        self.head_dim: int = config.head_dim
        self.num_kv_heads: int = getattr(config, "num_key_value_heads", self.num_heads)
        if self.num_heads % max(1, self.num_kv_heads) != 0:
            raise ValueError(
                f"num_heads ({self.num_heads}) must be divisible by "
                f"num_kv_heads ({self.num_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
        self.local_window: int = config.local_window
        self.rank: int = config.rank
        self.kernel_dim: int = config.kernel_feature_dim
        self.dropout = nn.Dropout(config.dropout)
        self.is_decoder: bool = getattr(config, "is_decoder", False)
        self.is_nested_encoder: bool = False

        self.q_proj = FactorisedLinear(self.hidden_size, self.all_head_size, rank=self.rank)
        self.k_proj = FactorisedLinear(self.hidden_size, self.all_kv_head_size, rank=self.rank)
        self.v_proj = FactorisedLinear(self.hidden_size, self.all_kv_head_size, rank=self.rank)
        self.out_proj = FactorisedLinear(self.all_head_size, self.hidden_size, rank=self.rank)

        if config.use_adaptive_pruning:
            self.prune_gate: nn.Sequential | None = nn.Sequential(
                nn.Linear(self.hidden_size, max(1, self.hidden_size // 4)),
                nn.ReLU(),
                nn.Linear(max(1, self.hidden_size // 4), self.num_heads),
            )
        else:
            self.prune_gate = None

        self.fusion_gate = nn.Linear(self.hidden_size, self.num_heads)
        self.fusion_gate.bias.data.fill_(config.gate_init_bias)

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

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

    @property
    def q_proj2(self):
        return self.q_proj.up

    @property
    def k_proj1(self):
        return self.k_proj.down

    @property
    def k_proj2(self):
        return self.k_proj.up

    @property
    def v_proj1(self):
        return self.v_proj.down

    @property
    def v_proj2(self):
        return self.v_proj.up

    @property
    def out_proj1(self):
        return self.out_proj.down

    @property
    def out_proj2(self):
        return self.out_proj.up

    @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

    # ------------------------------------------------------------------
    # Weight transfer (same logic as V5)
    # ------------------------------------------------------------------
    def _copy_pretrained(self, old_attn: nn.Module) -> None:
        self.copy_weights(old_attn)

    def copy_weights(self, old_attn: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

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

        # GPT-2 fused c_attn (Conv1D)
        if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            c_attn_mod = old_attn.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                is_conv1d = (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                )
                if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                    weight = weight.t().contiguous()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bd = bias.data if bias is not None else None
                split = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split, :],
                    bd[:split] if bd is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split, :],
                    bd[split : 2 * split] if bd is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :, :],
                    bd[2 * split :] if bd is not None else None,
                )

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

        # BERT / RoBERTa nested encoder
        if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        cls_name = old_attn.__class__.__name__
        has_out_proj = any(
            hasattr(old_attn, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
            )
            or "RobertaSelfAttention" in str(type(old_attn))
            or "BertSelfAttention" in str(type(old_attn))
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        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:
            raise RuntimeError(
                f"AegisAttentionV2.copy_weights: missing {missing} from "
                f"{old_attn.__class__.__name__}. Refusing random init."
            )

        for proj_key in ("q", "k", "v"):
            w = extracted[proj_key][0]
            if w is not None:
                norm = float(torch.linalg.norm(w.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "AegisAttentionV2: weight for '%s' has near-zero norm (%.2e).",
                        proj_key,
                        norm,
                    )

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

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    target.decompose_from_weight(_orient(w, in_dim), b)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
            self.is_decoder = True

    # ------------------------------------------------------------------
    # V2-FIX-07: safe from_distilbert_pretrained
    # ------------------------------------------------------------------
    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> "AegisAttentionV2":
        if model_config is None:
            out_mod = getattr(old_attn, "out_lin", None) or getattr(old_attn, "dense", None)
            if out_mod is None or not hasattr(out_mod, "out_features"):
                raise ValueError(
                    "from_distilbert_pretrained: could not infer hidden_size "
                    "from out_lin or dense. Pass model_config explicitly."
                )

            class DummyConfig:
                pass

            model_config = DummyConfig()
            model_config.hidden_size = out_mod.out_features
            n_heads = getattr(
                old_attn,
                "n_heads",
                getattr(
                    old_attn,
                    "num_attention_heads",
                    getattr(old_attn, "num_heads", None),
                ),
            )
            if n_heads is None:
                n_heads = max(1, model_config.hidden_size // 64)
            model_config.num_attention_heads = n_heads
            model_config.num_key_value_heads = n_heads
            model_config.head_dim = model_config.hidden_size // n_heads
        return cls(config=model_config, pretrained_attn=old_attn, **kwargs)

    # ------------------------------------------------------------------
    # V2-FIX-05: unified mask handling
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """Return boolean (B, N_kv): True = key is valid. None if no mask."""
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            # (B, 1, N_q, N_kv) -> OR over query dim -> (B, N_kv)
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # V2-FIX-02 + V2-FIX-09: 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        Hkv = self.num_kv_heads
        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]

        # ---- Mask (V2-FIX-05) ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv_new, device)

        # ---- Projections ----
        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
        V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

        # ---- RoPE with past offset ----
        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_rotary_pos_emb(Q, K_new, cos, sin)

        # ---- KV-cache concat ----
        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(getattr(self, "layer_idx", 0), len(past_key_value.key_cache) - 1)
                past_k, past_v = (
                    past_key_value.key_cache[idx],
                    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:
            idx = getattr(self, "layer_idx", 0)
            past_key_value.update(K_new, V_new, idx)
        else:
            pass

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

        # ---- GQA expansion ----
        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        # ---- Local attention (V2-FIX-03) ----
        local_out = _local_window_attention_v2(
            Q,
            K_full,
            V_full,
            window=self.local_window,
            token_valid=token_valid,
            is_decoder=self.is_decoder and not is_cross,
            past_len=past_len,
        )  # (B, H, N, D)

        # ---- Global attention (V2-FIX-08) ----
        if not is_cross:
            if self.prune_gate is not None:
                keep = torch.sigmoid(self.prune_gate(hidden_states))  # (B, N, H)
                # V2-FIX-06: expand across GQA heads
                if self.kv_repeat > 1:
                    keep_kv = keep[:, :, : self.num_kv_heads]
                else:
                    keep_kv = keep
                Kg = K_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
                Vg = V_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
            else:
                Kg, Vg = K_cat, V_cat

            if self.kv_repeat > 1:
                Kg = repeat_kv(Kg, self.kv_repeat)
                Vg = repeat_kv(Vg, self.kv_repeat)

            global_out = _global_linear_attention_v2(
                Q,
                Kg,
                Vg,
                token_valid=token_valid,
                is_decoder=self.is_decoder,
            )  # (B, H, N, D)

            alpha = torch.sigmoid(self.fusion_gate(hidden_states)).unsqueeze(-1)  # (B, N, H, 1)
            alpha = alpha.transpose(1, 2)  # (B, H, N, 1)
            fused = alpha * local_out + (1.0 - alpha) * global_out
        else:
            # Cross-attention: local only (global over encoder is questionable)
            fused = local_out

        # ---- head_mask (V2-FIX-02) ----
        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)

        # ---- Reshape to (B, N, H*D) ----
        out = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)
            output = self.dropout(output)
        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)
        return output.to(hidden_states.dtype), None

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.zeros((), device=self.out_proj.up.weight.device)

Every 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.

Static methods

def from_distilbert_pretrained(old_attn: torch.nn.modules.module.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> AegisAttentionV2

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 k_proj1

Expand source code

@property
def k_proj1(self):
    return self.k_proj.down

prop k_proj2

Expand source code

@property
def k_proj2(self):
    return self.k_proj.up

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 out_proj1

Expand source code

@property
def out_proj1(self):
    return self.out_proj.down

prop out_proj2

Expand source code

@property
def out_proj2(self):
    return self.out_proj.up

prop q_lin

Expand source code

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

prop q_proj1

Expand source code

@property
def q_proj1(self):
    return self.q_proj.down

prop q_proj2

Expand source code

@property
def q_proj2(self):
    return self.q_proj.up

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 v_proj1

Expand source code

@property
def v_proj1(self):
    return self.v_proj.down

prop v_proj2

Expand source code

@property
def v_proj2(self):
    return self.v_proj.up

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:
    from neuralops.core.adapters import UniversalAdapter

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

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

    # GPT-2 fused c_attn (Conv1D)
    if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        c_attn_mod = old_attn.c_attn
        base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
        weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
        if weight is not None:
            weight = weight.data.clone()
            is_conv1d = (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or hasattr(c_attn_mod, "nf")
                or hasattr(base_c, "nf")
            )
            if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                weight = weight.t().contiguous()
            bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            bd = bias.data if bias is not None else None
            split = weight.shape[0] // 3
            extracted["q"] = (
                weight[:split, :],
                bd[:split] if bd is not None else None,
            )
            extracted["k"] = (
                weight[split : 2 * split, :],
                bd[split : 2 * split] if bd is not None else None,
            )
            extracted["v"] = (
                weight[2 * split :, :],
                bd[2 * split :] if bd is not None else None,
            )

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

    # BERT / RoBERTa nested encoder
    if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        s = old_attn.self
        if hasattr(s, "query"):
            extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
        if hasattr(s, "key"):
            extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
        if hasattr(s, "value"):
            extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

    named_map = {
        "q": ["q_proj", "q_lin", "query"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    for proj, names in named_map.items():
        if extracted[proj][0] is None:
            for name in names:
                if hasattr(old_attn, name):
                    layer = getattr(old_attn, name)
                    if hasattr(layer, "weight"):
                        extracted[proj] = (
                            layer.weight.data,
                            getattr(layer, "bias", None),
                        )
                        break

    cls_name = old_attn.__class__.__name__
    has_out_proj = any(
        hasattr(old_attn, a)
        for a in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    if (
        hasattr(old_attn, "out_lin")
        or hasattr(old_attn, "c_proj")
        or hasattr(old_attn, "o_proj")
    ):
        self.is_nested_encoder = False
    elif (
        hasattr(old_attn, "self")
        and hasattr(old_attn.self, "query")
        and hasattr(old_attn, "output")
        and hasattr(old_attn.output, "dense")
    ) or (
        cls_name
        in (
            "BertSelfAttention",
            "RobertaSelfAttention",
            "XLMRobertaSelfAttention",
            "ElectraSelfAttention",
        )
        or "RobertaSelfAttention" in str(type(old_attn))
        or "BertSelfAttention" in str(type(old_attn))
        or (cls_name.endswith("SelfAttention") and not has_out_proj)
    ):
        self.is_nested_encoder = True
    else:
        self.is_nested_encoder = False

    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:
        raise RuntimeError(
            f"AegisAttentionV2.copy_weights: missing {missing} from "
            f"{old_attn.__class__.__name__}. Refusing random init."
        )

    for proj_key in ("q", "k", "v"):
        w = extracted[proj_key][0]
        if w is not None:
            norm = float(torch.linalg.norm(w.float()).item())
            if norm < 1e-4:
                logger.warning(
                    "AegisAttentionV2: weight for '%s' has near-zero norm (%.2e).",
                    proj_key,
                    norm,
                )

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

    with torch.no_grad():
        for proj_key, target in (
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ):
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                target.decompose_from_weight(_orient(w, in_dim), b)
            elif proj_key == "o" and self.is_nested_encoder:
                eye_w = torch.eye(
                    self.all_head_size,
                    self.all_head_size,
                    device=self.out_proj.up.weight.device,
                    dtype=self.out_proj.up.weight.dtype,
                )
                target.decompose_from_weight(eye_w, None)

    if hasattr(old_attn, "is_decoder"):
        self.is_decoder = bool(old_attn.is_decoder)
    elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
        self.is_decoder = True

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    H, D = self.num_heads, self.head_dim
    Hkv = self.num_kv_heads
    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]

    # ---- Mask (V2-FIX-05) ----
    eff_mask = encoder_attention_mask if is_cross else attention_mask
    token_valid = self._prepare_valid_mask(eff_mask, B, N_kv_new, device)

    # ---- Projections ----
    Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
    K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
    V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

    # ---- RoPE with past offset ----
    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_rotary_pos_emb(Q, K_new, cos, sin)

    # ---- KV-cache concat ----
    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(getattr(self, "layer_idx", 0), len(past_key_value.key_cache) - 1)
            past_k, past_v = (
                past_key_value.key_cache[idx],
                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:
        idx = getattr(self, "layer_idx", 0)
        past_key_value.update(K_new, V_new, idx)
    else:
        pass

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

    # ---- GQA expansion ----
    K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
    V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

    # ---- Local attention (V2-FIX-03) ----
    local_out = _local_window_attention_v2(
        Q,
        K_full,
        V_full,
        window=self.local_window,
        token_valid=token_valid,
        is_decoder=self.is_decoder and not is_cross,
        past_len=past_len,
    )  # (B, H, N, D)

    # ---- Global attention (V2-FIX-08) ----
    if not is_cross:
        if self.prune_gate is not None:
            keep = torch.sigmoid(self.prune_gate(hidden_states))  # (B, N, H)
            # V2-FIX-06: expand across GQA heads
            if self.kv_repeat > 1:
                keep_kv = keep[:, :, : self.num_kv_heads]
            else:
                keep_kv = keep
            Kg = K_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
            Vg = V_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
        else:
            Kg, Vg = K_cat, V_cat

        if self.kv_repeat > 1:
            Kg = repeat_kv(Kg, self.kv_repeat)
            Vg = repeat_kv(Vg, self.kv_repeat)

        global_out = _global_linear_attention_v2(
            Q,
            Kg,
            Vg,
            token_valid=token_valid,
            is_decoder=self.is_decoder,
        )  # (B, H, N, D)

        alpha = torch.sigmoid(self.fusion_gate(hidden_states)).unsqueeze(-1)  # (B, N, H, 1)
        alpha = alpha.transpose(1, 2)  # (B, H, N, 1)
        fused = alpha * local_out + (1.0 - alpha) * global_out
    else:
        # Cross-attention: local only (global over encoder is questionable)
        fused = local_out

    # ---- head_mask (V2-FIX-02) ----
    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)

    # ---- Reshape to (B, N, H*D) ----
    out = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
    out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

    if self.is_nested_encoder:
        output = out
    else:
        output = self.out_proj(out)
        output = self.dropout(output)
    output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)
    return output.to(hidden_states.dtype), None

Must return (output, attention_weights_or_None).

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.zeros((), device=self.out_proj.up.weight.device)

class AegisAttentionV1 (config: Any = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("aegis_v2", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v1", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v3", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v4", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v1_module", supports_causal=True, supports_gqa=True)
class AegisAttentionV2(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        if config is None or not isinstance(config, AegisConfigV2):
            h_size = (
                kwargs.get(
                    "hidden_size",
                    getattr(
                        config,
                        "hidden_size",
                        getattr(config, "d_model", getattr(config, "dim", 768)),
                    ),
                )
                if config is not None
                else 768
            )
            n_heads = (
                kwargs.get(
                    "num_heads",
                    getattr(config, "num_attention_heads", getattr(config, "num_heads", 12)),
                )
                if config is not None
                else 12
            )
            h_dim = (
                kwargs.get(
                    "head_dim",
                    getattr(
                        config,
                        "head_dim",
                        getattr(config, "d_kv", h_size // max(1, n_heads)),
                    ),
                )
                if config is not None
                else h_size // max(1, n_heads)
            )
            n_kv_heads = (
                kwargs.get(
                    "num_key_value_heads",
                    getattr(
                        config,
                        "num_key_value_heads",
                        getattr(config, "num_kv_heads", n_heads),
                    ),
                )
                if config is not None
                else n_heads
            )

            # V2-FIX-01: robust decoder detection
            dec = _detect_decoder(config, is_decoder, pretrained_attn)

            config = AegisConfigV2(
                hidden_size=h_size,
                num_attention_heads=n_heads,
                num_heads=n_heads,
                num_key_value_heads=n_kv_heads,
                head_dim=h_dim,
                local_window=(
                    kwargs.get(
                        "local_window",
                        getattr(config, "block_size", getattr(config, "local_window", 64)),
                    )
                    if config is not None
                    else 64
                ),
                rank=(
                    kwargs.pop("rank", getattr(config, "lora_r", getattr(config, "rank", 64)))
                    if config is not None
                    else 64
                ),
                kernel_feature_dim=(
                    getattr(
                        config,
                        "rff_features",
                        getattr(config, "kernel_feature_dim", 128),
                    )
                    if config is not None
                    else 128
                ),
                dropout=(
                    kwargs.get("dropout", getattr(config, "dropout", 0.1))
                    if config is not None
                    else 0.1
                ),
                use_adaptive_pruning=(
                    getattr(config, "use_gate", True) if config is not None else True
                ),
                gate_init_bias=(
                    getattr(config, "gate_init_bias", 1.8) if config is not None else 1.8
                ),
                is_decoder=dec,
            )

        super().__init__(config)
        self.config = config
        self.hidden_size: int = config.hidden_size
        self.num_heads: int = config.num_attention_heads
        self.head_dim: int = config.head_dim
        self.num_kv_heads: int = getattr(config, "num_key_value_heads", self.num_heads)
        if self.num_heads % max(1, self.num_kv_heads) != 0:
            raise ValueError(
                f"num_heads ({self.num_heads}) must be divisible by "
                f"num_kv_heads ({self.num_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
        self.local_window: int = config.local_window
        self.rank: int = config.rank
        self.kernel_dim: int = config.kernel_feature_dim
        self.dropout = nn.Dropout(config.dropout)
        self.is_decoder: bool = getattr(config, "is_decoder", False)
        self.is_nested_encoder: bool = False

        self.q_proj = FactorisedLinear(self.hidden_size, self.all_head_size, rank=self.rank)
        self.k_proj = FactorisedLinear(self.hidden_size, self.all_kv_head_size, rank=self.rank)
        self.v_proj = FactorisedLinear(self.hidden_size, self.all_kv_head_size, rank=self.rank)
        self.out_proj = FactorisedLinear(self.all_head_size, self.hidden_size, rank=self.rank)

        if config.use_adaptive_pruning:
            self.prune_gate: nn.Sequential | None = nn.Sequential(
                nn.Linear(self.hidden_size, max(1, self.hidden_size // 4)),
                nn.ReLU(),
                nn.Linear(max(1, self.hidden_size // 4), self.num_heads),
            )
        else:
            self.prune_gate = None

        self.fusion_gate = nn.Linear(self.hidden_size, self.num_heads)
        self.fusion_gate.bias.data.fill_(config.gate_init_bias)

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

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

    @property
    def q_proj2(self):
        return self.q_proj.up

    @property
    def k_proj1(self):
        return self.k_proj.down

    @property
    def k_proj2(self):
        return self.k_proj.up

    @property
    def v_proj1(self):
        return self.v_proj.down

    @property
    def v_proj2(self):
        return self.v_proj.up

    @property
    def out_proj1(self):
        return self.out_proj.down

    @property
    def out_proj2(self):
        return self.out_proj.up

    @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

    # ------------------------------------------------------------------
    # Weight transfer (same logic as V5)
    # ------------------------------------------------------------------
    def _copy_pretrained(self, old_attn: nn.Module) -> None:
        self.copy_weights(old_attn)

    def copy_weights(self, old_attn: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

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

        # GPT-2 fused c_attn (Conv1D)
        if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            c_attn_mod = old_attn.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                is_conv1d = (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                )
                if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                    weight = weight.t().contiguous()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bd = bias.data if bias is not None else None
                split = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split, :],
                    bd[:split] if bd is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split, :],
                    bd[split : 2 * split] if bd is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :, :],
                    bd[2 * split :] if bd is not None else None,
                )

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

        # BERT / RoBERTa nested encoder
        if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        cls_name = old_attn.__class__.__name__
        has_out_proj = any(
            hasattr(old_attn, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
            )
            or "RobertaSelfAttention" in str(type(old_attn))
            or "BertSelfAttention" in str(type(old_attn))
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        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:
            raise RuntimeError(
                f"AegisAttentionV2.copy_weights: missing {missing} from "
                f"{old_attn.__class__.__name__}. Refusing random init."
            )

        for proj_key in ("q", "k", "v"):
            w = extracted[proj_key][0]
            if w is not None:
                norm = float(torch.linalg.norm(w.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "AegisAttentionV2: weight for '%s' has near-zero norm (%.2e).",
                        proj_key,
                        norm,
                    )

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

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    target.decompose_from_weight(_orient(w, in_dim), b)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
            self.is_decoder = True

    # ------------------------------------------------------------------
    # V2-FIX-07: safe from_distilbert_pretrained
    # ------------------------------------------------------------------
    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> "AegisAttentionV2":
        if model_config is None:
            out_mod = getattr(old_attn, "out_lin", None) or getattr(old_attn, "dense", None)
            if out_mod is None or not hasattr(out_mod, "out_features"):
                raise ValueError(
                    "from_distilbert_pretrained: could not infer hidden_size "
                    "from out_lin or dense. Pass model_config explicitly."
                )

            class DummyConfig:
                pass

            model_config = DummyConfig()
            model_config.hidden_size = out_mod.out_features
            n_heads = getattr(
                old_attn,
                "n_heads",
                getattr(
                    old_attn,
                    "num_attention_heads",
                    getattr(old_attn, "num_heads", None),
                ),
            )
            if n_heads is None:
                n_heads = max(1, model_config.hidden_size // 64)
            model_config.num_attention_heads = n_heads
            model_config.num_key_value_heads = n_heads
            model_config.head_dim = model_config.hidden_size // n_heads
        return cls(config=model_config, pretrained_attn=old_attn, **kwargs)

    # ------------------------------------------------------------------
    # V2-FIX-05: unified mask handling
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """Return boolean (B, N_kv): True = key is valid. None if no mask."""
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            # (B, 1, N_q, N_kv) -> OR over query dim -> (B, N_kv)
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # V2-FIX-02 + V2-FIX-09: 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        Hkv = self.num_kv_heads
        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]

        # ---- Mask (V2-FIX-05) ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv_new, device)

        # ---- Projections ----
        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
        V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

        # ---- RoPE with past offset ----
        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_rotary_pos_emb(Q, K_new, cos, sin)

        # ---- KV-cache concat ----
        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(getattr(self, "layer_idx", 0), len(past_key_value.key_cache) - 1)
                past_k, past_v = (
                    past_key_value.key_cache[idx],
                    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:
            idx = getattr(self, "layer_idx", 0)
            past_key_value.update(K_new, V_new, idx)
        else:
            pass

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

        # ---- GQA expansion ----
        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        # ---- Local attention (V2-FIX-03) ----
        local_out = _local_window_attention_v2(
            Q,
            K_full,
            V_full,
            window=self.local_window,
            token_valid=token_valid,
            is_decoder=self.is_decoder and not is_cross,
            past_len=past_len,
        )  # (B, H, N, D)

        # ---- Global attention (V2-FIX-08) ----
        if not is_cross:
            if self.prune_gate is not None:
                keep = torch.sigmoid(self.prune_gate(hidden_states))  # (B, N, H)
                # V2-FIX-06: expand across GQA heads
                if self.kv_repeat > 1:
                    keep_kv = keep[:, :, : self.num_kv_heads]
                else:
                    keep_kv = keep
                Kg = K_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
                Vg = V_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
            else:
                Kg, Vg = K_cat, V_cat

            if self.kv_repeat > 1:
                Kg = repeat_kv(Kg, self.kv_repeat)
                Vg = repeat_kv(Vg, self.kv_repeat)

            global_out = _global_linear_attention_v2(
                Q,
                Kg,
                Vg,
                token_valid=token_valid,
                is_decoder=self.is_decoder,
            )  # (B, H, N, D)

            alpha = torch.sigmoid(self.fusion_gate(hidden_states)).unsqueeze(-1)  # (B, N, H, 1)
            alpha = alpha.transpose(1, 2)  # (B, H, N, 1)
            fused = alpha * local_out + (1.0 - alpha) * global_out
        else:
            # Cross-attention: local only (global over encoder is questionable)
            fused = local_out

        # ---- head_mask (V2-FIX-02) ----
        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)

        # ---- Reshape to (B, N, H*D) ----
        out = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)
            output = self.dropout(output)
        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)
        return output.to(hidden_states.dtype), None

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.zeros((), device=self.out_proj.up.weight.device)

Every 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.

Static methods

def from_distilbert_pretrained(old_attn: torch.nn.modules.module.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> AegisAttentionV2

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 k_proj1

Expand source code

@property
def k_proj1(self):
    return self.k_proj.down

prop k_proj2

Expand source code

@property
def k_proj2(self):
    return self.k_proj.up

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 out_proj1

Expand source code

@property
def out_proj1(self):
    return self.out_proj.down

prop out_proj2

Expand source code

@property
def out_proj2(self):
    return self.out_proj.up

prop q_lin

Expand source code

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

prop q_proj1

Expand source code

@property
def q_proj1(self):
    return self.q_proj.down

prop q_proj2

Expand source code

@property
def q_proj2(self):
    return self.q_proj.up

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 v_proj1

Expand source code

@property
def v_proj1(self):
    return self.v_proj.down

prop v_proj2

Expand source code

@property
def v_proj2(self):
    return self.v_proj.up

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:
    from neuralops.core.adapters import UniversalAdapter

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

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

    # GPT-2 fused c_attn (Conv1D)
    if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        c_attn_mod = old_attn.c_attn
        base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
        weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
        if weight is not None:
            weight = weight.data.clone()
            is_conv1d = (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or hasattr(c_attn_mod, "nf")
                or hasattr(base_c, "nf")
            )
            if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                weight = weight.t().contiguous()
            bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            bd = bias.data if bias is not None else None
            split = weight.shape[0] // 3
            extracted["q"] = (
                weight[:split, :],
                bd[:split] if bd is not None else None,
            )
            extracted["k"] = (
                weight[split : 2 * split, :],
                bd[split : 2 * split] if bd is not None else None,
            )
            extracted["v"] = (
                weight[2 * split :, :],
                bd[2 * split :] if bd is not None else None,
            )

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

    # BERT / RoBERTa nested encoder
    if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        s = old_attn.self
        if hasattr(s, "query"):
            extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
        if hasattr(s, "key"):
            extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
        if hasattr(s, "value"):
            extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

    named_map = {
        "q": ["q_proj", "q_lin", "query"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    for proj, names in named_map.items():
        if extracted[proj][0] is None:
            for name in names:
                if hasattr(old_attn, name):
                    layer = getattr(old_attn, name)
                    if hasattr(layer, "weight"):
                        extracted[proj] = (
                            layer.weight.data,
                            getattr(layer, "bias", None),
                        )
                        break

    cls_name = old_attn.__class__.__name__
    has_out_proj = any(
        hasattr(old_attn, a)
        for a in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    if (
        hasattr(old_attn, "out_lin")
        or hasattr(old_attn, "c_proj")
        or hasattr(old_attn, "o_proj")
    ):
        self.is_nested_encoder = False
    elif (
        hasattr(old_attn, "self")
        and hasattr(old_attn.self, "query")
        and hasattr(old_attn, "output")
        and hasattr(old_attn.output, "dense")
    ) or (
        cls_name
        in (
            "BertSelfAttention",
            "RobertaSelfAttention",
            "XLMRobertaSelfAttention",
            "ElectraSelfAttention",
        )
        or "RobertaSelfAttention" in str(type(old_attn))
        or "BertSelfAttention" in str(type(old_attn))
        or (cls_name.endswith("SelfAttention") and not has_out_proj)
    ):
        self.is_nested_encoder = True
    else:
        self.is_nested_encoder = False

    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:
        raise RuntimeError(
            f"AegisAttentionV2.copy_weights: missing {missing} from "
            f"{old_attn.__class__.__name__}. Refusing random init."
        )

    for proj_key in ("q", "k", "v"):
        w = extracted[proj_key][0]
        if w is not None:
            norm = float(torch.linalg.norm(w.float()).item())
            if norm < 1e-4:
                logger.warning(
                    "AegisAttentionV2: weight for '%s' has near-zero norm (%.2e).",
                    proj_key,
                    norm,
                )

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

    with torch.no_grad():
        for proj_key, target in (
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ):
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                target.decompose_from_weight(_orient(w, in_dim), b)
            elif proj_key == "o" and self.is_nested_encoder:
                eye_w = torch.eye(
                    self.all_head_size,
                    self.all_head_size,
                    device=self.out_proj.up.weight.device,
                    dtype=self.out_proj.up.weight.dtype,
                )
                target.decompose_from_weight(eye_w, None)

    if hasattr(old_attn, "is_decoder"):
        self.is_decoder = bool(old_attn.is_decoder)
    elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
        self.is_decoder = True

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    H, D = self.num_heads, self.head_dim
    Hkv = self.num_kv_heads
    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]

    # ---- Mask (V2-FIX-05) ----
    eff_mask = encoder_attention_mask if is_cross else attention_mask
    token_valid = self._prepare_valid_mask(eff_mask, B, N_kv_new, device)

    # ---- Projections ----
    Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
    K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
    V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

    # ---- RoPE with past offset ----
    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_rotary_pos_emb(Q, K_new, cos, sin)

    # ---- KV-cache concat ----
    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(getattr(self, "layer_idx", 0), len(past_key_value.key_cache) - 1)
            past_k, past_v = (
                past_key_value.key_cache[idx],
                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:
        idx = getattr(self, "layer_idx", 0)
        past_key_value.update(K_new, V_new, idx)
    else:
        pass

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

    # ---- GQA expansion ----
    K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
    V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

    # ---- Local attention (V2-FIX-03) ----
    local_out = _local_window_attention_v2(
        Q,
        K_full,
        V_full,
        window=self.local_window,
        token_valid=token_valid,
        is_decoder=self.is_decoder and not is_cross,
        past_len=past_len,
    )  # (B, H, N, D)

    # ---- Global attention (V2-FIX-08) ----
    if not is_cross:
        if self.prune_gate is not None:
            keep = torch.sigmoid(self.prune_gate(hidden_states))  # (B, N, H)
            # V2-FIX-06: expand across GQA heads
            if self.kv_repeat > 1:
                keep_kv = keep[:, :, : self.num_kv_heads]
            else:
                keep_kv = keep
            Kg = K_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
            Vg = V_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
        else:
            Kg, Vg = K_cat, V_cat

        if self.kv_repeat > 1:
            Kg = repeat_kv(Kg, self.kv_repeat)
            Vg = repeat_kv(Vg, self.kv_repeat)

        global_out = _global_linear_attention_v2(
            Q,
            Kg,
            Vg,
            token_valid=token_valid,
            is_decoder=self.is_decoder,
        )  # (B, H, N, D)

        alpha = torch.sigmoid(self.fusion_gate(hidden_states)).unsqueeze(-1)  # (B, N, H, 1)
        alpha = alpha.transpose(1, 2)  # (B, H, N, 1)
        fused = alpha * local_out + (1.0 - alpha) * global_out
    else:
        # Cross-attention: local only (global over encoder is questionable)
        fused = local_out

    # ---- head_mask (V2-FIX-02) ----
    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)

    # ---- Reshape to (B, N, H*D) ----
    out = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
    out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

    if self.is_nested_encoder:
        output = out
    else:
        output = self.out_proj(out)
        output = self.dropout(output)
    output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)
    return output.to(hidden_states.dtype), None

Must return (output, attention_weights_or_None).

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.zeros((), device=self.out_proj.up.weight.device)

class AegisAttentionV2 (config: Any = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("aegis_v2", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v1", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v3", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v4", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v1_module", supports_causal=True, supports_gqa=True)
class AegisAttentionV2(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        if config is None or not isinstance(config, AegisConfigV2):
            h_size = (
                kwargs.get(
                    "hidden_size",
                    getattr(
                        config,
                        "hidden_size",
                        getattr(config, "d_model", getattr(config, "dim", 768)),
                    ),
                )
                if config is not None
                else 768
            )
            n_heads = (
                kwargs.get(
                    "num_heads",
                    getattr(config, "num_attention_heads", getattr(config, "num_heads", 12)),
                )
                if config is not None
                else 12
            )
            h_dim = (
                kwargs.get(
                    "head_dim",
                    getattr(
                        config,
                        "head_dim",
                        getattr(config, "d_kv", h_size // max(1, n_heads)),
                    ),
                )
                if config is not None
                else h_size // max(1, n_heads)
            )
            n_kv_heads = (
                kwargs.get(
                    "num_key_value_heads",
                    getattr(
                        config,
                        "num_key_value_heads",
                        getattr(config, "num_kv_heads", n_heads),
                    ),
                )
                if config is not None
                else n_heads
            )

            # V2-FIX-01: robust decoder detection
            dec = _detect_decoder(config, is_decoder, pretrained_attn)

            config = AegisConfigV2(
                hidden_size=h_size,
                num_attention_heads=n_heads,
                num_heads=n_heads,
                num_key_value_heads=n_kv_heads,
                head_dim=h_dim,
                local_window=(
                    kwargs.get(
                        "local_window",
                        getattr(config, "block_size", getattr(config, "local_window", 64)),
                    )
                    if config is not None
                    else 64
                ),
                rank=(
                    kwargs.pop("rank", getattr(config, "lora_r", getattr(config, "rank", 64)))
                    if config is not None
                    else 64
                ),
                kernel_feature_dim=(
                    getattr(
                        config,
                        "rff_features",
                        getattr(config, "kernel_feature_dim", 128),
                    )
                    if config is not None
                    else 128
                ),
                dropout=(
                    kwargs.get("dropout", getattr(config, "dropout", 0.1))
                    if config is not None
                    else 0.1
                ),
                use_adaptive_pruning=(
                    getattr(config, "use_gate", True) if config is not None else True
                ),
                gate_init_bias=(
                    getattr(config, "gate_init_bias", 1.8) if config is not None else 1.8
                ),
                is_decoder=dec,
            )

        super().__init__(config)
        self.config = config
        self.hidden_size: int = config.hidden_size
        self.num_heads: int = config.num_attention_heads
        self.head_dim: int = config.head_dim
        self.num_kv_heads: int = getattr(config, "num_key_value_heads", self.num_heads)
        if self.num_heads % max(1, self.num_kv_heads) != 0:
            raise ValueError(
                f"num_heads ({self.num_heads}) must be divisible by "
                f"num_kv_heads ({self.num_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
        self.local_window: int = config.local_window
        self.rank: int = config.rank
        self.kernel_dim: int = config.kernel_feature_dim
        self.dropout = nn.Dropout(config.dropout)
        self.is_decoder: bool = getattr(config, "is_decoder", False)
        self.is_nested_encoder: bool = False

        self.q_proj = FactorisedLinear(self.hidden_size, self.all_head_size, rank=self.rank)
        self.k_proj = FactorisedLinear(self.hidden_size, self.all_kv_head_size, rank=self.rank)
        self.v_proj = FactorisedLinear(self.hidden_size, self.all_kv_head_size, rank=self.rank)
        self.out_proj = FactorisedLinear(self.all_head_size, self.hidden_size, rank=self.rank)

        if config.use_adaptive_pruning:
            self.prune_gate: nn.Sequential | None = nn.Sequential(
                nn.Linear(self.hidden_size, max(1, self.hidden_size // 4)),
                nn.ReLU(),
                nn.Linear(max(1, self.hidden_size // 4), self.num_heads),
            )
        else:
            self.prune_gate = None

        self.fusion_gate = nn.Linear(self.hidden_size, self.num_heads)
        self.fusion_gate.bias.data.fill_(config.gate_init_bias)

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

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

    @property
    def q_proj2(self):
        return self.q_proj.up

    @property
    def k_proj1(self):
        return self.k_proj.down

    @property
    def k_proj2(self):
        return self.k_proj.up

    @property
    def v_proj1(self):
        return self.v_proj.down

    @property
    def v_proj2(self):
        return self.v_proj.up

    @property
    def out_proj1(self):
        return self.out_proj.down

    @property
    def out_proj2(self):
        return self.out_proj.up

    @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

    # ------------------------------------------------------------------
    # Weight transfer (same logic as V5)
    # ------------------------------------------------------------------
    def _copy_pretrained(self, old_attn: nn.Module) -> None:
        self.copy_weights(old_attn)

    def copy_weights(self, old_attn: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

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

        # GPT-2 fused c_attn (Conv1D)
        if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            c_attn_mod = old_attn.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                is_conv1d = (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                )
                if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                    weight = weight.t().contiguous()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bd = bias.data if bias is not None else None
                split = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split, :],
                    bd[:split] if bd is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split, :],
                    bd[split : 2 * split] if bd is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :, :],
                    bd[2 * split :] if bd is not None else None,
                )

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

        # BERT / RoBERTa nested encoder
        if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        cls_name = old_attn.__class__.__name__
        has_out_proj = any(
            hasattr(old_attn, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
            )
            or "RobertaSelfAttention" in str(type(old_attn))
            or "BertSelfAttention" in str(type(old_attn))
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        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:
            raise RuntimeError(
                f"AegisAttentionV2.copy_weights: missing {missing} from "
                f"{old_attn.__class__.__name__}. Refusing random init."
            )

        for proj_key in ("q", "k", "v"):
            w = extracted[proj_key][0]
            if w is not None:
                norm = float(torch.linalg.norm(w.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "AegisAttentionV2: weight for '%s' has near-zero norm (%.2e).",
                        proj_key,
                        norm,
                    )

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

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    target.decompose_from_weight(_orient(w, in_dim), b)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
            self.is_decoder = True

    # ------------------------------------------------------------------
    # V2-FIX-07: safe from_distilbert_pretrained
    # ------------------------------------------------------------------
    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> "AegisAttentionV2":
        if model_config is None:
            out_mod = getattr(old_attn, "out_lin", None) or getattr(old_attn, "dense", None)
            if out_mod is None or not hasattr(out_mod, "out_features"):
                raise ValueError(
                    "from_distilbert_pretrained: could not infer hidden_size "
                    "from out_lin or dense. Pass model_config explicitly."
                )

            class DummyConfig:
                pass

            model_config = DummyConfig()
            model_config.hidden_size = out_mod.out_features
            n_heads = getattr(
                old_attn,
                "n_heads",
                getattr(
                    old_attn,
                    "num_attention_heads",
                    getattr(old_attn, "num_heads", None),
                ),
            )
            if n_heads is None:
                n_heads = max(1, model_config.hidden_size // 64)
            model_config.num_attention_heads = n_heads
            model_config.num_key_value_heads = n_heads
            model_config.head_dim = model_config.hidden_size // n_heads
        return cls(config=model_config, pretrained_attn=old_attn, **kwargs)

    # ------------------------------------------------------------------
    # V2-FIX-05: unified mask handling
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """Return boolean (B, N_kv): True = key is valid. None if no mask."""
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            # (B, 1, N_q, N_kv) -> OR over query dim -> (B, N_kv)
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # V2-FIX-02 + V2-FIX-09: 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        Hkv = self.num_kv_heads
        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]

        # ---- Mask (V2-FIX-05) ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv_new, device)

        # ---- Projections ----
        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
        V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

        # ---- RoPE with past offset ----
        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_rotary_pos_emb(Q, K_new, cos, sin)

        # ---- KV-cache concat ----
        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(getattr(self, "layer_idx", 0), len(past_key_value.key_cache) - 1)
                past_k, past_v = (
                    past_key_value.key_cache[idx],
                    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:
            idx = getattr(self, "layer_idx", 0)
            past_key_value.update(K_new, V_new, idx)
        else:
            pass

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

        # ---- GQA expansion ----
        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        # ---- Local attention (V2-FIX-03) ----
        local_out = _local_window_attention_v2(
            Q,
            K_full,
            V_full,
            window=self.local_window,
            token_valid=token_valid,
            is_decoder=self.is_decoder and not is_cross,
            past_len=past_len,
        )  # (B, H, N, D)

        # ---- Global attention (V2-FIX-08) ----
        if not is_cross:
            if self.prune_gate is not None:
                keep = torch.sigmoid(self.prune_gate(hidden_states))  # (B, N, H)
                # V2-FIX-06: expand across GQA heads
                if self.kv_repeat > 1:
                    keep_kv = keep[:, :, : self.num_kv_heads]
                else:
                    keep_kv = keep
                Kg = K_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
                Vg = V_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
            else:
                Kg, Vg = K_cat, V_cat

            if self.kv_repeat > 1:
                Kg = repeat_kv(Kg, self.kv_repeat)
                Vg = repeat_kv(Vg, self.kv_repeat)

            global_out = _global_linear_attention_v2(
                Q,
                Kg,
                Vg,
                token_valid=token_valid,
                is_decoder=self.is_decoder,
            )  # (B, H, N, D)

            alpha = torch.sigmoid(self.fusion_gate(hidden_states)).unsqueeze(-1)  # (B, N, H, 1)
            alpha = alpha.transpose(1, 2)  # (B, H, N, 1)
            fused = alpha * local_out + (1.0 - alpha) * global_out
        else:
            # Cross-attention: local only (global over encoder is questionable)
            fused = local_out

        # ---- head_mask (V2-FIX-02) ----
        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)

        # ---- Reshape to (B, N, H*D) ----
        out = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)
            output = self.dropout(output)
        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)
        return output.to(hidden_states.dtype), None

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.zeros((), device=self.out_proj.up.weight.device)

Every 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.

Static methods

def from_distilbert_pretrained(old_attn: torch.nn.modules.module.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> AegisAttentionV2

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 k_proj1

Expand source code

@property
def k_proj1(self):
    return self.k_proj.down

prop k_proj2

Expand source code

@property
def k_proj2(self):
    return self.k_proj.up

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 out_proj1

Expand source code

@property
def out_proj1(self):
    return self.out_proj.down

prop out_proj2

Expand source code

@property
def out_proj2(self):
    return self.out_proj.up

prop q_lin

Expand source code

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

prop q_proj1

Expand source code

@property
def q_proj1(self):
    return self.q_proj.down

prop q_proj2

Expand source code

@property
def q_proj2(self):
    return self.q_proj.up

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 v_proj1

Expand source code

@property
def v_proj1(self):
    return self.v_proj.down

prop v_proj2

Expand source code

@property
def v_proj2(self):
    return self.v_proj.up

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:
    from neuralops.core.adapters import UniversalAdapter

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

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

    # GPT-2 fused c_attn (Conv1D)
    if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        c_attn_mod = old_attn.c_attn
        base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
        weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
        if weight is not None:
            weight = weight.data.clone()
            is_conv1d = (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or hasattr(c_attn_mod, "nf")
                or hasattr(base_c, "nf")
            )
            if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                weight = weight.t().contiguous()
            bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            bd = bias.data if bias is not None else None
            split = weight.shape[0] // 3
            extracted["q"] = (
                weight[:split, :],
                bd[:split] if bd is not None else None,
            )
            extracted["k"] = (
                weight[split : 2 * split, :],
                bd[split : 2 * split] if bd is not None else None,
            )
            extracted["v"] = (
                weight[2 * split :, :],
                bd[2 * split :] if bd is not None else None,
            )

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

    # BERT / RoBERTa nested encoder
    if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        s = old_attn.self
        if hasattr(s, "query"):
            extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
        if hasattr(s, "key"):
            extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
        if hasattr(s, "value"):
            extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

    named_map = {
        "q": ["q_proj", "q_lin", "query"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    for proj, names in named_map.items():
        if extracted[proj][0] is None:
            for name in names:
                if hasattr(old_attn, name):
                    layer = getattr(old_attn, name)
                    if hasattr(layer, "weight"):
                        extracted[proj] = (
                            layer.weight.data,
                            getattr(layer, "bias", None),
                        )
                        break

    cls_name = old_attn.__class__.__name__
    has_out_proj = any(
        hasattr(old_attn, a)
        for a in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    if (
        hasattr(old_attn, "out_lin")
        or hasattr(old_attn, "c_proj")
        or hasattr(old_attn, "o_proj")
    ):
        self.is_nested_encoder = False
    elif (
        hasattr(old_attn, "self")
        and hasattr(old_attn.self, "query")
        and hasattr(old_attn, "output")
        and hasattr(old_attn.output, "dense")
    ) or (
        cls_name
        in (
            "BertSelfAttention",
            "RobertaSelfAttention",
            "XLMRobertaSelfAttention",
            "ElectraSelfAttention",
        )
        or "RobertaSelfAttention" in str(type(old_attn))
        or "BertSelfAttention" in str(type(old_attn))
        or (cls_name.endswith("SelfAttention") and not has_out_proj)
    ):
        self.is_nested_encoder = True
    else:
        self.is_nested_encoder = False

    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:
        raise RuntimeError(
            f"AegisAttentionV2.copy_weights: missing {missing} from "
            f"{old_attn.__class__.__name__}. Refusing random init."
        )

    for proj_key in ("q", "k", "v"):
        w = extracted[proj_key][0]
        if w is not None:
            norm = float(torch.linalg.norm(w.float()).item())
            if norm < 1e-4:
                logger.warning(
                    "AegisAttentionV2: weight for '%s' has near-zero norm (%.2e).",
                    proj_key,
                    norm,
                )

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

    with torch.no_grad():
        for proj_key, target in (
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ):
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                target.decompose_from_weight(_orient(w, in_dim), b)
            elif proj_key == "o" and self.is_nested_encoder:
                eye_w = torch.eye(
                    self.all_head_size,
                    self.all_head_size,
                    device=self.out_proj.up.weight.device,
                    dtype=self.out_proj.up.weight.dtype,
                )
                target.decompose_from_weight(eye_w, None)

    if hasattr(old_attn, "is_decoder"):
        self.is_decoder = bool(old_attn.is_decoder)
    elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
        self.is_decoder = True

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    H, D = self.num_heads, self.head_dim
    Hkv = self.num_kv_heads
    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]

    # ---- Mask (V2-FIX-05) ----
    eff_mask = encoder_attention_mask if is_cross else attention_mask
    token_valid = self._prepare_valid_mask(eff_mask, B, N_kv_new, device)

    # ---- Projections ----
    Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
    K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
    V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

    # ---- RoPE with past offset ----
    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_rotary_pos_emb(Q, K_new, cos, sin)

    # ---- KV-cache concat ----
    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(getattr(self, "layer_idx", 0), len(past_key_value.key_cache) - 1)
            past_k, past_v = (
                past_key_value.key_cache[idx],
                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:
        idx = getattr(self, "layer_idx", 0)
        past_key_value.update(K_new, V_new, idx)
    else:
        pass

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

    # ---- GQA expansion ----
    K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
    V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

    # ---- Local attention (V2-FIX-03) ----
    local_out = _local_window_attention_v2(
        Q,
        K_full,
        V_full,
        window=self.local_window,
        token_valid=token_valid,
        is_decoder=self.is_decoder and not is_cross,
        past_len=past_len,
    )  # (B, H, N, D)

    # ---- Global attention (V2-FIX-08) ----
    if not is_cross:
        if self.prune_gate is not None:
            keep = torch.sigmoid(self.prune_gate(hidden_states))  # (B, N, H)
            # V2-FIX-06: expand across GQA heads
            if self.kv_repeat > 1:
                keep_kv = keep[:, :, : self.num_kv_heads]
            else:
                keep_kv = keep
            Kg = K_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
            Vg = V_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
        else:
            Kg, Vg = K_cat, V_cat

        if self.kv_repeat > 1:
            Kg = repeat_kv(Kg, self.kv_repeat)
            Vg = repeat_kv(Vg, self.kv_repeat)

        global_out = _global_linear_attention_v2(
            Q,
            Kg,
            Vg,
            token_valid=token_valid,
            is_decoder=self.is_decoder,
        )  # (B, H, N, D)

        alpha = torch.sigmoid(self.fusion_gate(hidden_states)).unsqueeze(-1)  # (B, N, H, 1)
        alpha = alpha.transpose(1, 2)  # (B, H, N, 1)
        fused = alpha * local_out + (1.0 - alpha) * global_out
    else:
        # Cross-attention: local only (global over encoder is questionable)
        fused = local_out

    # ---- head_mask (V2-FIX-02) ----
    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)

    # ---- Reshape to (B, N, H*D) ----
    out = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
    out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

    if self.is_nested_encoder:
        output = out
    else:
        output = self.out_proj(out)
        output = self.dropout(output)
    output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)
    return output.to(hidden_states.dtype), None

Must return (output, attention_weights_or_None).

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.zeros((), device=self.out_proj.up.weight.device)

class AegisAttentionV4 (config: Any = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("aegis_v2", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v1", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v3", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v4", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v1_module", supports_causal=True, supports_gqa=True)
class AegisAttentionV2(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        if config is None or not isinstance(config, AegisConfigV2):
            h_size = (
                kwargs.get(
                    "hidden_size",
                    getattr(
                        config,
                        "hidden_size",
                        getattr(config, "d_model", getattr(config, "dim", 768)),
                    ),
                )
                if config is not None
                else 768
            )
            n_heads = (
                kwargs.get(
                    "num_heads",
                    getattr(config, "num_attention_heads", getattr(config, "num_heads", 12)),
                )
                if config is not None
                else 12
            )
            h_dim = (
                kwargs.get(
                    "head_dim",
                    getattr(
                        config,
                        "head_dim",
                        getattr(config, "d_kv", h_size // max(1, n_heads)),
                    ),
                )
                if config is not None
                else h_size // max(1, n_heads)
            )
            n_kv_heads = (
                kwargs.get(
                    "num_key_value_heads",
                    getattr(
                        config,
                        "num_key_value_heads",
                        getattr(config, "num_kv_heads", n_heads),
                    ),
                )
                if config is not None
                else n_heads
            )

            # V2-FIX-01: robust decoder detection
            dec = _detect_decoder(config, is_decoder, pretrained_attn)

            config = AegisConfigV2(
                hidden_size=h_size,
                num_attention_heads=n_heads,
                num_heads=n_heads,
                num_key_value_heads=n_kv_heads,
                head_dim=h_dim,
                local_window=(
                    kwargs.get(
                        "local_window",
                        getattr(config, "block_size", getattr(config, "local_window", 64)),
                    )
                    if config is not None
                    else 64
                ),
                rank=(
                    kwargs.pop("rank", getattr(config, "lora_r", getattr(config, "rank", 64)))
                    if config is not None
                    else 64
                ),
                kernel_feature_dim=(
                    getattr(
                        config,
                        "rff_features",
                        getattr(config, "kernel_feature_dim", 128),
                    )
                    if config is not None
                    else 128
                ),
                dropout=(
                    kwargs.get("dropout", getattr(config, "dropout", 0.1))
                    if config is not None
                    else 0.1
                ),
                use_adaptive_pruning=(
                    getattr(config, "use_gate", True) if config is not None else True
                ),
                gate_init_bias=(
                    getattr(config, "gate_init_bias", 1.8) if config is not None else 1.8
                ),
                is_decoder=dec,
            )

        super().__init__(config)
        self.config = config
        self.hidden_size: int = config.hidden_size
        self.num_heads: int = config.num_attention_heads
        self.head_dim: int = config.head_dim
        self.num_kv_heads: int = getattr(config, "num_key_value_heads", self.num_heads)
        if self.num_heads % max(1, self.num_kv_heads) != 0:
            raise ValueError(
                f"num_heads ({self.num_heads}) must be divisible by "
                f"num_kv_heads ({self.num_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
        self.local_window: int = config.local_window
        self.rank: int = config.rank
        self.kernel_dim: int = config.kernel_feature_dim
        self.dropout = nn.Dropout(config.dropout)
        self.is_decoder: bool = getattr(config, "is_decoder", False)
        self.is_nested_encoder: bool = False

        self.q_proj = FactorisedLinear(self.hidden_size, self.all_head_size, rank=self.rank)
        self.k_proj = FactorisedLinear(self.hidden_size, self.all_kv_head_size, rank=self.rank)
        self.v_proj = FactorisedLinear(self.hidden_size, self.all_kv_head_size, rank=self.rank)
        self.out_proj = FactorisedLinear(self.all_head_size, self.hidden_size, rank=self.rank)

        if config.use_adaptive_pruning:
            self.prune_gate: nn.Sequential | None = nn.Sequential(
                nn.Linear(self.hidden_size, max(1, self.hidden_size // 4)),
                nn.ReLU(),
                nn.Linear(max(1, self.hidden_size // 4), self.num_heads),
            )
        else:
            self.prune_gate = None

        self.fusion_gate = nn.Linear(self.hidden_size, self.num_heads)
        self.fusion_gate.bias.data.fill_(config.gate_init_bias)

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

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

    @property
    def q_proj2(self):
        return self.q_proj.up

    @property
    def k_proj1(self):
        return self.k_proj.down

    @property
    def k_proj2(self):
        return self.k_proj.up

    @property
    def v_proj1(self):
        return self.v_proj.down

    @property
    def v_proj2(self):
        return self.v_proj.up

    @property
    def out_proj1(self):
        return self.out_proj.down

    @property
    def out_proj2(self):
        return self.out_proj.up

    @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

    # ------------------------------------------------------------------
    # Weight transfer (same logic as V5)
    # ------------------------------------------------------------------
    def _copy_pretrained(self, old_attn: nn.Module) -> None:
        self.copy_weights(old_attn)

    def copy_weights(self, old_attn: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

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

        # GPT-2 fused c_attn (Conv1D)
        if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            c_attn_mod = old_attn.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                is_conv1d = (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                )
                if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                    weight = weight.t().contiguous()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bd = bias.data if bias is not None else None
                split = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split, :],
                    bd[:split] if bd is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split, :],
                    bd[split : 2 * split] if bd is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :, :],
                    bd[2 * split :] if bd is not None else None,
                )

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

        # BERT / RoBERTa nested encoder
        if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        cls_name = old_attn.__class__.__name__
        has_out_proj = any(
            hasattr(old_attn, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
            )
            or "RobertaSelfAttention" in str(type(old_attn))
            or "BertSelfAttention" in str(type(old_attn))
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        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:
            raise RuntimeError(
                f"AegisAttentionV2.copy_weights: missing {missing} from "
                f"{old_attn.__class__.__name__}. Refusing random init."
            )

        for proj_key in ("q", "k", "v"):
            w = extracted[proj_key][0]
            if w is not None:
                norm = float(torch.linalg.norm(w.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "AegisAttentionV2: weight for '%s' has near-zero norm (%.2e).",
                        proj_key,
                        norm,
                    )

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

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    target.decompose_from_weight(_orient(w, in_dim), b)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
            self.is_decoder = True

    # ------------------------------------------------------------------
    # V2-FIX-07: safe from_distilbert_pretrained
    # ------------------------------------------------------------------
    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> "AegisAttentionV2":
        if model_config is None:
            out_mod = getattr(old_attn, "out_lin", None) or getattr(old_attn, "dense", None)
            if out_mod is None or not hasattr(out_mod, "out_features"):
                raise ValueError(
                    "from_distilbert_pretrained: could not infer hidden_size "
                    "from out_lin or dense. Pass model_config explicitly."
                )

            class DummyConfig:
                pass

            model_config = DummyConfig()
            model_config.hidden_size = out_mod.out_features
            n_heads = getattr(
                old_attn,
                "n_heads",
                getattr(
                    old_attn,
                    "num_attention_heads",
                    getattr(old_attn, "num_heads", None),
                ),
            )
            if n_heads is None:
                n_heads = max(1, model_config.hidden_size // 64)
            model_config.num_attention_heads = n_heads
            model_config.num_key_value_heads = n_heads
            model_config.head_dim = model_config.hidden_size // n_heads
        return cls(config=model_config, pretrained_attn=old_attn, **kwargs)

    # ------------------------------------------------------------------
    # V2-FIX-05: unified mask handling
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """Return boolean (B, N_kv): True = key is valid. None if no mask."""
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            # (B, 1, N_q, N_kv) -> OR over query dim -> (B, N_kv)
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # V2-FIX-02 + V2-FIX-09: 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        Hkv = self.num_kv_heads
        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]

        # ---- Mask (V2-FIX-05) ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv_new, device)

        # ---- Projections ----
        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K_new = self.k_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()
        V_new = self.v_proj(kv_input).view(B, N_kv_new, Hkv, D).transpose(1, 2).contiguous()

        # ---- RoPE with past offset ----
        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_rotary_pos_emb(Q, K_new, cos, sin)

        # ---- KV-cache concat ----
        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(getattr(self, "layer_idx", 0), len(past_key_value.key_cache) - 1)
                past_k, past_v = (
                    past_key_value.key_cache[idx],
                    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:
            idx = getattr(self, "layer_idx", 0)
            past_key_value.update(K_new, V_new, idx)
        else:
            pass

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

        # ---- GQA expansion ----
        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        # ---- Local attention (V2-FIX-03) ----
        local_out = _local_window_attention_v2(
            Q,
            K_full,
            V_full,
            window=self.local_window,
            token_valid=token_valid,
            is_decoder=self.is_decoder and not is_cross,
            past_len=past_len,
        )  # (B, H, N, D)

        # ---- Global attention (V2-FIX-08) ----
        if not is_cross:
            if self.prune_gate is not None:
                keep = torch.sigmoid(self.prune_gate(hidden_states))  # (B, N, H)
                # V2-FIX-06: expand across GQA heads
                if self.kv_repeat > 1:
                    keep_kv = keep[:, :, : self.num_kv_heads]
                else:
                    keep_kv = keep
                Kg = K_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
                Vg = V_cat * keep_kv.transpose(1, 2).unsqueeze(-1)
            else:
                Kg, Vg = K_cat, V_cat

            if self.kv_repeat > 1:
                Kg = repeat_kv(Kg, self.kv_repeat)
                Vg = repeat_kv(Vg, self.kv_repeat)

            global_out = _global_linear_attention_v2(
                Q,
                Kg,
                Vg,
                token_valid=token_valid,
                is_decoder=self.is_decoder,
            )  # (B, H, N, D)

            alpha = torch.sigmoid(self.fusion_gate(hidden_states)).unsqueeze(-1)  # (B, N, H, 1)
            alpha = alpha.transpose(1, 2)  # (B, H, N, 1)
            fused = alpha * local_out + (1.0 - alpha) * global_out
        else:
            # Cross-attention: local only (global over encoder is questionable)
            fused = local_out

        # ---- head_mask (V2-FIX-02) ----
        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)

        # ---- Reshape to (B, N, H*D) ----
        out = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)
            output = self.dropout(output)
        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)
        return output.to(hidden_states.dtype), None

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.zeros((), device=self.out_proj.up.weight.device)

Every 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.

Static methods

def from_distilbert_pretrained(old_attn: torch.nn.modules.module.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> AegisAttentionV2

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 k_proj1

Expand source code

@property
def k_proj1(self):
    return self.k_proj.down

prop k_proj2

Expand source code

@property
def k_proj2(self):
    return self.k_proj.up

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 out_proj1

Expand source code

@property
def out_proj1(self):
    return self.out_proj.down

prop out_proj2

Expand source code

@property
def out_proj2(self):
    return self.out_proj.up

prop q_lin

Expand source code

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

prop q_proj1

Expand source code

@property
def q_proj1(self):
    return self.q_proj.down

prop q_proj2

Expand source code

@property
def q_proj2(self):
    return self.q_proj.up

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 v_proj1

Expand source code

@property
def v_proj1(self):
    return self.v_proj.down

prop v_proj2

Expand source code

@property
def v_proj2(self):
    return self.v_proj.up

prop value

Expand source code

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

Methods

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.zeros((), device=self.out_proj.up.weight.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AegisAttentionV6 (config=None, pretrained_qkv=None, pretrained_attn=None, is_decoder=None, **kwargs)

Expand source code

@ATTENTION_REGISTRY.register("aegis_v6", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_v6_attn", supports_causal=True, supports_gqa=True)
class AegisAttentionV3(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config=None,
        pretrained_qkv=None,
        pretrained_attn=None,
        is_decoder=None,
        **kwargs,
    ):
        if not isinstance(config, AegisConfigV3):
            h_size = getattr(
                config,
                "hidden_size",
                getattr(
                    config,
                    "d_model",
                    getattr(config, "n_embd", getattr(config, "dim", None)),
                ),
            )
            if h_size is None and pretrained_attn is not None:
                for obj in [
                    pretrained_attn,
                    getattr(pretrained_attn, "self", None),
                    getattr(pretrained_attn, "output", None),
                ]:
                    if (
                        obj is not None
                        and hasattr(obj, "dense")
                        and hasattr(obj.dense, "out_features")
                    ):
                        h_size = obj.dense.out_features
                        break
            if h_size is None:
                h_size = kwargs.get("hidden_size", 768)
            n_heads = getattr(
                config,
                "num_attention_heads",
                getattr(
                    config,
                    "num_heads",
                    getattr(config, "n_heads", getattr(config, "n_head", None)),
                ),
            ) or kwargs.get("num_heads", 12)
            h_dim = getattr(config, "head_dim", getattr(config, "d_kv", None)) or kwargs.get(
                "head_dim", h_size // max(1, n_heads)
            )
            n_kv_heads = getattr(
                config, "num_key_value_heads", getattr(config, "num_kv_heads", None)
            ) or kwargs.get("num_key_value_heads", n_heads)
            dec = self._detect_decoder(config, is_decoder, pretrained_attn)
            config = AegisConfigV3(
                hidden_size=h_size,
                num_heads=n_heads,
                num_key_value_heads=n_kv_heads,
                head_dim=h_dim,
                window_size=getattr(config, "window_size", 64),
                svd_rank=getattr(config, "svd_rank", min(128, h_size) if h_size else 128),
                low_rank_update=getattr(config, "low_rank_update", 16),
                kernel_feature_dim=getattr(config, "kernel_feature_dim", 128),
                dropout=getattr(config, "dropout", 0.1),
                is_decoder=dec,
            )
        super().__init__(config)
        self.cfg = config
        self.hidden_size = config.hidden_size
        self.h = config.num_heads
        self.num_heads = config.num_heads
        self.num_kv_heads = getattr(config, "num_key_value_heads", self.h) or self.h
        if self.h % self.num_kv_heads != 0:
            raise ValueError(
                f"num_heads {self.h} not divisible by num_kv_heads {self.num_kv_heads}"
            )
        self.kv_repeat = max(1, self.h // max(1, self.num_kv_heads))
        self.d = config.head_dim
        self.head_dim = self.d
        self.all_head_size = self.h * self.d
        self.all_kv_head_size = self.num_kv_heads * self.d
        self.window = config.window_size
        self.is_decoder = getattr(config, "is_decoder", False)
        self.is_nested_encoder = False
        self.chunk_size = config.chunk_size

        self.q_proj = FactorisedLinear(
            self.hidden_size,
            self.all_head_size,
            config.svd_rank,
            config.low_rank_update,
        )
        self.k_proj = FactorisedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            config.svd_rank,
            config.low_rank_update,
        )
        self.v_proj = FactorisedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            config.svd_rank,
            config.low_rank_update,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            self.hidden_size,
            config.svd_rank,
            config.low_rank_update,
        )

        self.pre_norm = config.use_pre_norm
        if self.pre_norm:
            self.q_norm = nn.LayerNorm(self.d)
            self.k_norm = nn.LayerNorm(self.d)

        self.kernel_map = StableKernelFeatureMap(
            KernelFeatureConfig(
                head_dim=self.d,
                num_features=config.kernel_feature_dim,
                method=config.kernel_method,
                clamp_min=config.kernel_clamp_min,
                clamp_max=config.kernel_clamp_max,
                use_float64_exp=config.kernel_use_float64,
                seed=config.kernel_seed,
            )
        )

        self.gates = GatingModule(
            config.hidden_size,
            self.h,
            self.d,
            config.use_query_gate,
            config.use_input_gate,
        )
        self.fusion = ContextFusionGate(config.hidden_size, self.h, self.d, config.fusion_gate_bias)

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

    @staticmethod
    def _detect_decoder(config, is_decoder, old_attn):
        if is_decoder is not None:
            return bool(is_decoder)
        if config is not None:
            if getattr(config, "is_decoder", False) or getattr(config, "causal", False):
                return True
            mt = str(getattr(config, "model_type", "")).lower()
            if mt.startswith(
                (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "gemma",
                    "phi",
                    "opt",
                    "bloom",
                    "falcon",
                    "olmo",
                    "mpt",
                    "stablelm",
                )
            ):
                return True
        if old_attn is not None:
            if getattr(old_attn, "is_decoder", False):
                return True
            if hasattr(old_attn, "c_attn"):
                return True
            if "gpt" in old_attn.__class__.__name__.lower():
                return True
        return False

    @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

    # ------------------------------------------------------------------
    # Mask utilities
    # ------------------------------------------------------------------
    @staticmethod
    def _raw_valid(mask, B, N_kv, device):
        if mask is None:
            return torch.ones(B, 1, 1, N_kv, device=device, dtype=torch.bool)
        if mask.dim() == 2:
            if mask.dtype == torch.bool:
                v = mask
            elif mask.dtype.is_floating_point:
                v = mask > -1e4
            else:
                v = mask > 0
            return v[:, None, None, :].to(torch.bool)
        if mask.dim() == 4:
            if mask.dtype == torch.bool:
                v = mask
            elif mask.dtype.is_floating_point:
                v = mask > -1e4
            else:
                v = mask > 0
            return v.to(torch.bool)
        raise ValueError(f"Unsupported attention_mask dim: {mask.dim()}")

    def _prepare_valid_mask(
        self,
        attention_mask: torch.Tensor | None,
        B: int,
        N_q: int,
        N_kv: int,
        device: torch.device,
        is_cross: bool,
    ) -> torch.Tensor:
        """
        Build a boolean validity mask of shape (B, 1, N_q, N_kv).
        True = key is allowed for the corresponding query.

        Causal masking is applied ONLY when:
        - the module is a decoder, AND
        - this is self-attention (not cross-attention).

        Handles three decoding regimes correctly via absolute positions:
        * prefill:       N_q == N_kv   (full sequence, causal mask diag-aligned)
        * single decode: N_q == 1      (single query sees all past keys)
        * multi decode:  N_q > 1, N_q < N_kv (chunked decode with past cache)
        In all three cases the query at absolute position `past_len + i` can attend
        to keys at absolute positions `<= past_len + i`.
        """
        # 1) Raw validity from padding (before any causal logic)
        raw = self._raw_valid(attention_mask, B, N_kv, device)  # (B, 1, 1 or N_q, N_kv)

        # 2) Broadcast padding along query dim if needed
        if raw.shape[-2] == 1 and N_q != 1:
            raw = raw.expand(B, 1, N_q, N_kv)

        # 3) Causal mask (decoder self-attention only)
        if self.is_decoder and not is_cross:
            # Sanity check: shapes must line up
            if raw.shape[-2] != N_q or raw.shape[-1] != N_kv:
                raise ValueError(
                    f"_prepare_valid_mask: shape mismatch: raw={tuple(raw.shape)} "
                    f"expected (B, 1, {N_q}, {N_kv})"
                )

            # Absolute positions: query i has absolute index (past_len + i).
            # past_len = N_kv - N_q is exact for both prefill and decode-with-cache.
            past_len = N_kv - N_q
            idx_q = (past_len + torch.arange(N_q, device=device)).unsqueeze(1)  # (N_q, 1)
            idx_k = torch.arange(N_kv, device=device).unsqueeze(0)  # (1, N_kv)

            # Causal: query q can attend to key k iff abs_k <= abs_q
            causal = (idx_k <= idx_q).unsqueeze(0).unsqueeze(0)  # (1, 1, N_q, N_kv)
            raw = raw & causal

        return raw

    # ------------------------------------------------------------------
    # Sliding window — FIX-S1, FIX-S2
    # ------------------------------------------------------------------
    def _sliding_window_attention(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        valid: torch.Tensor,
        is_cross: bool,
        return_attn: bool = False,
    ):
        """
        Sliding-window + causal attention.

        Args:
            Q:       (B, h, N, d)
            K:       (B, h, N_kv, d)
            V:       (B, h, N_kv, d)
            valid:   (B, 1, N, N_kv) boolean — padding + causal pre-computed
                    by `_prepare_valid_mask`. This function additionally restricts
                    attention to a local band of width `self.window`.
            is_cross: whether this call is cross-attention (skip causal band).
            return_attn: if True, also return the attention matrix.

        Returns:
            (out, attn) if return_attn else (out, None)
        """
        B, h, N, d = Q.shape
        N_k = K.shape[2]
        W = self.window
        scale = d**-0.5

        # 1) Raw scores (full N x N_kv matrix; for production replace with a
        #    FlashAttention / Triton kernel to actually get O(N*W) memory).
        scores = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale

        # 2) Absolute positions
        #    past_len is exact for prefill (0), single decode (N_kv-1),
        #    and multi decode (N_kv-N).
        past_len = (N_k - N) if (self.is_decoder and not is_cross and N_k >= N) else 0
        abs_q = (past_len + torch.arange(N, device=Q.device)).unsqueeze(1)  # (N, 1)
        abs_k = torch.arange(N_k, device=Q.device).unsqueeze(0)  # (1, N_k)

        # 3) Local band: |abs_q - abs_k| <= W
        dist = (abs_q - abs_k).abs()  # (N, N_k)
        band = (dist <= W).unsqueeze(0).unsqueeze(0)  # (1, 1, N, N_k)

        # 4) Causal band: query at abs_q can attend to key at abs_k iff abs_k <= abs_q.
        #    Applied for ALL decoder self-attention cases (prefill, single decode,
        #    multi-token decode). It degenerates to "all keys allowed" when N == 1
        #    because abs_q = N_k-1 >= every abs_k.
        if self.is_decoder and not is_cross:
            causal = (abs_k <= abs_q).unsqueeze(0).unsqueeze(0)  # (1, 1, N, N_k)
            band = band & causal

        # 5) Apply band + caller-provided valid mask (padding + causal).
        mask_val = -1e4
        scores = scores.masked_fill(~band, mask_val)
        scores = scores.masked_fill(~valid, mask_val)
        scores = torch.nan_to_num(scores, nan=mask_val, posinf=1e4, neginf=mask_val)

        # 6) Safe softmax: rows with no valid key → zero output (not uniform).
        any_valid = valid.any(dim=-1, keepdim=True).to(scores.dtype)  # (B, 1, N, 1)
        attn = F.softmax(scores.float(), dim=-1).to(Q.dtype) * any_valid
        attn = F.dropout(attn, p=self.cfg.dropout, training=self.training)

        # 7) Weighted sum
        out = torch.matmul(attn, V)

        if return_attn:
            return out, attn
        return out, None

    # ------------------------------------------------------------------
    # Non-causal linear attention (encoder)
    # ------------------------------------------------------------------
    def _global_linear_attention(self, Q, K, V, valid_kv_only, q_gate, in_gate):
        Qk = self.q_norm(Q) if self.pre_norm else Q
        Kk = self.k_norm(K) if self.pre_norm else K
        Q_feat = self.kernel_map(Qk)
        K_feat = self.kernel_map(Kk)
        if q_gate is not None:
            Q_feat = Q_feat * q_gate
        if in_gate is not None and in_gate.shape[2] == K_feat.shape[2]:
            K_feat = K_feat * in_gate
            V = V * in_gate
        valid_k = valid_kv_only.transpose(-1, -2).to(K_feat.dtype)
        K_feat = K_feat * valid_k
        V = V * valid_k
        KV = torch.matmul(K_feat.transpose(-2, -1).contiguous(), V)
        Z = K_feat.sum(dim=-2, keepdim=True)
        num = torch.matmul(Q_feat, KV)
        den = torch.matmul(Q_feat, Z.transpose(-2, -1).contiguous()).clamp(min=1e-6)
        return torch.nan_to_num(num / den, nan=0.0)

        # ------------------------------------------------------------------
        # Causal linear attention  — FIX-C3, FIX-C4
        # ------------------------------------------------------------------

    def _causal_linear_attention(
        self, Q, K, V, valid, q_gate, in_gate, chunk_size: int | None = None
    ):
        """
        Chunked causal linear attention.

        Three cases:
        * N == N_k (training / prefill): full chunked cumsum with causal structure.
        * N == 1 (single-token decode with cache): exact, since the single query
            attends to all keys in the cache (causal mask trivially holds).
        * N > 1 and N != N_k (multi-token decode): FALLBACK to sliding-window
            attention over the full key history. This avoids the current
            non-causal approximation which is known to be inaccurate.
        """
        B, h, N, d = Q.shape
        N_k = K.shape[2]
        if chunk_size is None:
            chunk_size = self.chunk_size

        Qk = self.q_norm(Q) if self.pre_norm else Q
        Kk = self.k_norm(K) if self.pre_norm else K
        Q_feat = self.kernel_map(Qk)
        K_feat = self.kernel_map(Kk)

        if q_gate is not None:
            Q_feat = Q_feat * q_gate
        if in_gate is not None and in_gate.shape[2] == K_feat.shape[2]:
            K_feat = K_feat * in_gate
            V = V * in_gate

        valid_k = valid.any(dim=-2, keepdim=True).transpose(-1, -2).to(K_feat.dtype)
        K_feat = K_feat * valid_k
        V = V * valid_k

        # ---- Case: decode with cache (N != N_k) ----
        if N_k != N:
            # FIX: For single-token decode (N==1), non-causal over full history is EXACT.
            # For multi-token decode, fall back to exact sliding-window attention
            # rather than the inaccurate non-causal approximation.
            if N == 1:
                KV = torch.matmul(K_feat.transpose(-2, -1).contiguous(), V)
                Z = K_feat.sum(dim=-2, keepdim=True)
                num = torch.matmul(Q_feat, KV)
                den = torch.matmul(Q_feat, Z.transpose(-2, -1).contiguous()).clamp(min=1e-6)
                return torch.nan_to_num(num / den, nan=0.0)
            else:
                # Multi-token decode: use exact sliding-window attention
                logger.warning(
                    "_causal_linear_attention: multi-token decode (N=%d, N_k=%d). "
                    "Falling back to exact sliding-window attention.",
                    N,
                    N_k,
                )
                _fallback_out, _ = self._sliding_window_attention(
                    Q, K, V, valid, False, return_attn=False
                )
                return _fallback_out

        # ---- Case: N == N_k (training / full prefill) ----
        M = Q_feat.shape[-1]
        state_kv = torch.zeros(B, h, M, d, device=Q.device, dtype=Q_feat.dtype)
        state_k = torch.zeros(B, h, M, device=Q.device, dtype=Q_feat.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]

            # FIX-C3: correct outer product orientation
            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("bhnm,bhnmd->bhnd", Qc, KV_total)
            den = torch.einsum("bhnm,bhnm->bhn", Qc, K_total).clamp(min=1e-6)
            outs.append(num / den.unsqueeze(-1))

            state_kv = KV_total[:, :, -1]
            state_k = K_total[:, :, -1]

        out = torch.cat(outs, dim=2)
        return torch.nan_to_num(out, nan=0.0)

    # ------------------------------------------------------------------
    # KV-cache
    # ------------------------------------------------------------------
    @staticmethod
    def _read_past_kv(past_key_value, layer_idx=0):
        if past_key_value is None:
            return None, None
        if hasattr(past_key_value, "key_cache") and hasattr(past_key_value, "value_cache"):
            keys = past_key_value.key_cache
            vals = past_key_value.value_cache
            if len(keys) == 0:
                return None, None
            idx = min(layer_idx, len(keys) - 1)
            return keys[idx], vals[idx]
        if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
            return past_key_value[0], past_key_value[1]
        return None, None

    @staticmethod
    def _past_len(past_k):
        return 0 if past_k is None else past_k.shape[-2]

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
        position_embeddings=None,
        layer_idx=0,
        *args,
        **kwargs,
    ):
        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

        Q = self._reshape(self.q_proj(hidden_states), self.h)
        K_new = self._reshape(self.k_proj(kv_input), self.num_kv_heads)
        V_new = self._reshape(self.v_proj(kv_input), self.num_kv_heads)

        # RoPE with past offset
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            past_k_probe, _ = self._read_past_kv(past_key_value, layer_idx)
            past_len = self._past_len(past_k_probe)
            cos = cos[..., past_len : past_len + N, :]
            sin = sin[..., past_len : past_len + N, :]
            Q, K_new = apply_rotary_pos_emb(Q, K_new, cos, sin)

        past_k, past_v = self._read_past_kv(past_key_value, layer_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)
        else:
            pass

        N_kv = K_cat.shape[2]
        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        valid = self._prepare_valid_mask(eff_mask, B, N, N_kv, device, is_cross)

        q_gate, in_gate = self.gates(hidden_states, Q)

        if self.is_decoder:
            local_out, local_attn = self._sliding_window_attention(
                Q, K_full, V_full, valid, is_cross, return_attn=bool(output_attentions)
            )
            global_out = self._causal_linear_attention(Q, K_full, V_full, valid, q_gate, in_gate)
            fuse_gate = self.fusion(hidden_states, local_out, global_out)
            fused = fuse_gate * local_out + (1.0 - fuse_gate) * global_out
        else:
            local_out, local_attn = self._sliding_window_attention(
                Q, K_full, V_full, valid, is_cross, return_attn=bool(output_attentions)
            )
            valid_kv_only = valid.any(dim=-2, keepdim=True)
            global_out = self._global_linear_attention(
                Q, K_full, V_full, valid_kv_only, q_gate, in_gate
            )
            fuse_gate = self.fusion(hidden_states, local_out, global_out)
            fused = fuse_gate * local_out + (1.0 - fuse_gate) * global_out

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)
            fused = fused * head_mask.to(fused.dtype)

        fused = fused.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)
        fused = torch.nan_to_num(fused, nan=0.0, posinf=1e4, neginf=-1e4)
        out = fused if self.is_nested_encoder else self.out_proj(fused)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        attn_ret = local_attn if output_attentions else None
        # Always return a plain 2-tuple so `out, weights = attn(...)` unpacking works
        # for both encoder (legacy_bert) and decoder paths.
        return out, attn_ret

    def _reshape(self, x, num_heads):
        B, N, _ = x.shape
        return x.view(B, N, num_heads, self.d).permute(0, 2, 1, 3)

    def get_ortho_loss(self):
        return torch.zeros((), device=self.out_proj.bias.device)

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn):
        from neuralops.core.adapters import UniversalAdapter

        cls_name = old_attn.__class__.__name__
        has_out_proj_attr = any(
            hasattr(old_attn, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
            )
            or "RobertaSelfAttention" in str(type(old_attn))
            or "BertSelfAttention" in str(type(old_attn))
            or (cls_name.endswith("SelfAttention") and not has_out_proj_attr)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        extracted = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        for proj in ["q", "k", "v", "o"]:
            try:
                w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
                if w is not None:
                    extracted[proj] = (w, b)
            except Exception as e:
                logger.debug(f"UniversalAdapter failed for {proj}: {e}")

        if hasattr(old_attn, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            c_attn_mod = old_attn.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                is_conv1d = (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                )
                if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                    weight = weight.t().contiguous()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bd = bias.data if bias is not None else None
                split = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split, :],
                    bd[:split] if bd is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split, :],
                    bd[split : 2 * split] if bd is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :, :],
                    bd[2 * split :] if bd is not None else None,
                )

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

        if hasattr(old_attn, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

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

        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        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:
            raise RuntimeError(
                f"AegisAttentionV3.copy_weights: could not extract mandatory projections "
                f"{missing} from {old_attn.__class__.__name__}. "
                f"Refusing to install randomly-initialised weights."
            )

        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

        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    target.decompose_from_weight(_orient(w, in_dim), b)
                    # Warn if weight norm is suspiciously near-zero
                    norm = float(torch.linalg.norm(w.float()).item())
                    if norm < 1e-4:
                        logger.warning(
                            "AegisAttentionV3.copy_weights: extracted weight for '%s' has "
                            "near-zero Frobenius norm (%.2e) — weights may be uninitialised.",
                            proj_key,
                            norm,
                        )
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        device=self.out_proj.U.device,
                        dtype=self.out_proj.U.dtype,
                    )
                    target.decompose_from_weight(eye_w, None)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
            self.is_decoder = True

Every 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.

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 get_ortho_loss(self)

Expand source code

def get_ortho_loss(self):
    return torch.zeros((), device=self.out_proj.bias.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AegisConfigV3 (**kwargs: Any)

Expand source code

class AegisConfigV3:
    def __init__(self, **kwargs: Any):
        self.hidden_size = kwargs.get("hidden_size")
        self.num_heads = kwargs.get("num_heads", kwargs.get("num_attention_heads"))
        self.num_attention_heads = self.num_heads
        self.num_key_value_heads = kwargs.get("num_key_value_heads", self.num_heads)
        self.head_dim = kwargs.get("head_dim")
        if self.head_dim is None and self.hidden_size is not None and self.num_heads is not None:
            self.head_dim = self.hidden_size // max(1, self.num_heads)
        self.window_size = kwargs.get("window_size", 64)
        default_svd = min(384, self.hidden_size) if self.hidden_size is not None else 384
        self.svd_rank = kwargs.get("svd_rank", default_svd)
        self.low_rank_update = kwargs.get("low_rank_update", 16)
        self.kernel_feature_dim = kwargs.get("kernel_feature_dim", 128)
        self.dropout = kwargs.get("dropout", 0.1)
        self.use_query_gate = kwargs.get("use_query_gate", True)
        self.use_input_gate = kwargs.get("use_input_gate", True)
        self.use_pre_norm = kwargs.get("use_pre_norm", False)
        self.long_context_threshold = kwargs.get("long_context_threshold", 2048)
        self.max_global_tokens = kwargs.get("max_global_tokens", 1024)
        self.kernel_method = kwargs.get("kernel_method", "orthogonal")
        self.kernel_clamp_min = kwargs.get("kernel_clamp_min", -20.0)
        self.kernel_clamp_max = kwargs.get("kernel_clamp_max", 8.0)
        self.kernel_use_float64 = kwargs.get("kernel_use_float64", False)
        self.kernel_seed = kwargs.get("kernel_seed")
        self.fusion_gate_bias = kwargs.get("fusion_gate_bias", 1.8)
        self.is_decoder = kwargs.get("is_decoder", False)
        self.chunk_size = kwargs.get("chunk_size", 64)

class AegisConfigV5 (svd_energy: float = 0.95,
min_rank: int = 4,
max_rank: int = 384,
num_local_heads: int = 8,
num_global_heads: int = 4,
local_window: int = 64,
softpick_temperature: float = 1.0,
rff_features: int = 128,
use_zero_sum_keys: bool = True,
use_adaptive_gate: bool = True,
gate_init_bias: float = 1.8,
dropout: float = 0.1,
is_decoder: bool = False,
rff_trainable: bool = False,
use_sdpa_fast_path: bool = True,
causal_chunk_size: int = 64)

Expand source code

@dataclass
class AegisConfigV5:
    """Configuration for AegisAttention V5."""

    svd_energy: float = 0.95
    min_rank: int = 4
    max_rank: int = 384
    num_local_heads: int = 8
    num_global_heads: int = 4
    # V5-FIX-09: local_window = full width of the sliding window.
    # Effective radius = local_window // 2 (queried keys are |q - k| <= radius).
    local_window: int = 64
    softpick_temperature: float = 1.0
    rff_features: int = 128
    use_zero_sum_keys: bool = True
    use_adaptive_gate: bool = True
    gate_init_bias: float = 1.8
    dropout: float = 0.1
    is_decoder: bool = False
    # V5-FIX-14: rff_proj as buffer (fixed random features). Set True to make trainable.
    rff_trainable: bool = False
    # SDPA fast path for very short sequences.
    use_sdpa_fast_path: bool = True
    # Chunk size for causal kernel linear attention (lower = less memory, slower).
    causal_chunk_size: int = 64

Configuration for AegisAttention V5.

Instance variables

var causal_chunk_size : int

The type of the None singleton.

var dropout : float

The type of the None singleton.

var gate_init_bias : float

The type of the None singleton.

var is_decoder : bool

The type of the None singleton.

var local_window : int

The type of the None singleton.

var max_rank : int

The type of the None singleton.

var min_rank : int

The type of the None singleton.

var num_global_heads : int

The type of the None singleton.

var num_local_heads : int

The type of the None singleton.

var rff_features : int

The type of the None singleton.

var rff_trainable : bool

The type of the None singleton.

var softpick_temperature : float

The type of the None singleton.

var svd_energy : float

The type of the None singleton.

var use_adaptive_gate : bool

The type of the None singleton.

var use_sdpa_fast_path : bool

The type of the None singleton.

var use_zero_sum_keys : bool

The type of the None singleton.

class AegisHybridAttention (config: Any | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
pretrained_attn: nn.Module | None = None,
rank: int | None = None,
lora_rank: int | None = None,
is_decoder: bool | None = None,
whitening: dict[str, torch.Tensor] | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "aegis_hybrid", supports_causal=True, supports_gqa=True, supports_rope=True
)
@ATTENTION_REGISTRY.register(
    "aegis_hybrid_attn", supports_causal=True, supports_gqa=True, supports_rope=True
)
@ATTENTION_REGISTRY.register(
    "aegis_hybrid_attention",
    supports_causal=True,
    supports_gqa=True,
    supports_rope=True,
)
class AegisHybridAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        pretrained_attn: nn.Module | None = None,
        rank: int | None = None,
        lora_rank: int | None = None,
        is_decoder: bool | None = None,
        whitening: dict[str, torch.Tensor] | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank if rank is not None else 128)
        except TypeError:
            super().__init__(config)

        if not isinstance(config, AegisHybridConfig):
            cfg_src = config if config is not None else kwargs
            h_size = hidden_size or getattr(cfg_src, "hidden_size", 768)
            n_heads = num_heads or getattr(
                cfg_src, "num_attention_heads", getattr(cfg_src, "n_heads", 12)
            )
            n_kv = getattr(
                cfg_src,
                "num_key_value_heads",
                getattr(cfg_src, "num_kv_heads", n_heads),
            )
            h_dim = getattr(cfg_src, "head_dim", None) or (h_size // max(1, n_heads))
            r = rank or getattr(cfg_src, "rank", min(128, h_size))
            l_r = lora_rank or getattr(cfg_src, "lora_rank", 32)

            dec = _detect_decoder(config, is_decoder, pretrained_attn)
            config = AegisHybridConfig(
                hidden_size=h_size,
                num_attention_heads=n_heads,
                head_dim=h_dim,
                num_key_value_heads=n_kv,
                local_window=kwargs.get("local_window", 256),
                rank=r,
                lora_rank=l_r,
                is_decoder=dec,
                use_global=kwargs.get("use_global", True),
                kernel_feature_dim=kwargs.get("kernel_feature_dim", 128),
                chunk_size=kwargs.get("chunk_size", 64),
            )

        self.config = config
        self.hidden_size: int = int(config.hidden_size)
        self.num_heads: int = int(config.num_attention_heads)
        self.head_dim: int = int(config.head_dim)
        self.num_kv_heads: int = int(config.num_key_value_heads or self.num_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
        self.local_window: int = int(config.local_window)
        self.rank: int = max(1, min(int(config.rank), self.hidden_size))
        self.lora_rank: int = max(1, min(int(config.lora_rank), self.hidden_size))
        self.dropout = nn.Dropout(config.dropout)

        self.is_decoder: bool = _detect_decoder(config, is_decoder, pretrained_attn)
        self.use_global: bool = bool(config.use_global)
        self.is_nested_encoder: bool = False
        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

        self.whitening = whitening or {}

        # ---- projections ----
        self.q_proj = FactorizedLinear(
            self.hidden_size,
            self.all_head_size,
            rank=self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.k_proj = FactorizedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            rank=self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.v_proj = FactorizedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            rank=self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.out_proj = FactorizedLinear(
            self.all_head_size,
            self.hidden_size,
            rank=self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )

        # ---- kernel feature map ----
        self.kernel_map = KernelFeatureMap(
            head_dim=self.head_dim,
            num_features=config.kernel_feature_dim,
            seed=config.kernel_seed,
        )

        # ---- fusion gate (per-head, init ~0.88 to local) ----
        self.fusion_beta = nn.Parameter(torch.zeros(self.num_heads))

        # AH-FIX-12: storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 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

    # ------------------------------------------------------------------
    # AH-FIX-04: GQA-aware SVD init
    # ------------------------------------------------------------------
    @staticmethod
    def _svd_init(
        w: torch.Tensor,
        b: torch.Tensor | None,
        proj: FactorizedLinear,
        whitening_L: torch.Tensor | None = None,
    ) -> None:
        with torch.no_grad():
            W = w.detach().float()
            out_dim, in_dim = W.shape
            target_out, target_in = proj.out_features, proj.in_features

            # GQA: pool or repeat along output dim
            if out_dim > target_out and out_dim % target_out == 0:
                rep = out_dim // target_out
                W = W.view(target_out, rep, in_dim).mean(dim=1)
                if b is not None:
                    b = b.view(target_out, rep).mean(dim=1)
            elif out_dim < target_out and target_out % out_dim == 0:
                rep = target_out // out_dim
                W = W.repeat_interleave(rep, dim=0)
                if b is not None:
                    b = b.repeat_interleave(rep)

            if W.shape[0] != target_out or W.shape[1] != target_in:
                raise ValueError(
                    f"FactorizedLinear shape mismatch after GQA pooling: "
                    f"{tuple(W.shape)} vs ({target_out}, {target_in})"
                )

            # Optional whitening
            L = None
            if whitening_L is not None:
                Lw = whitening_L.float().to(W.device)
                if Lw.shape == (target_in, target_in):
                    eps = 1e-6 * torch.trace(Lw @ Lw.t()).item() / max(1, target_in)
                    L = Lw + eps * torch.eye(target_in, device=W.device, dtype=W.dtype)

            if L is not None:
                try:
                    W_eff = W @ L
                    U, S, Vh = torch.linalg.svd(W_eff, full_matrices=False)
                    L_inv = torch.linalg.inv(L)
                    Vh = Vh @ L_inv
                except Exception:
                    U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            else:
                U, S, Vh = torch.linalg.svd(W, full_matrices=False)

            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.detach()[:n].to(proj.up.bias.dtype))
            elif proj.up.bias is not None:
                proj.up.bias.data.zero_()

            # Invalidate cache
            proj._cached_weight = None
            proj._cached_versions = ()

    # ------------------------------------------------------------------
    # AH-FIX-05: copy_weights with proper Conv1D handling
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: nn.Module) -> None:
        try:
            for p in old_attn.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # UniversalAdapter first
        for p in ("q", "k", "v", "o"):
            try:
                w, b = UniversalAdapter.get_weight_bias(old_attn, p)
                if w is not None:
                    extracted[p] = (w, b)
            except Exception:
                pass

        # GPT-2 fused c_attn
        if hasattr(old_attn, "c_attn") and extracted["q"][0] is None:
            c = old_attn.c_attn
            c_w = c.weight.data.clone()
            # Conv1D stores weight as (in, out) → transpose to (out, in)
            if (
                c_w.shape[1] == 3 * self.hidden_size
                or c_w.ndim == 2
                and c_w.shape[0] < c_w.shape[1]
            ):
                c_w = c_w.t().contiguous()
            split = c_w.size(0) // 3
            c_b = getattr(c, "bias", None)
            c_b = c_b.data.clone() if c_b is not None else None
            extracted["q"] = (c_w[:split], c_b[:split] if c_b is not None else None)
            extracted["k"] = (
                c_w[split : 2 * split],
                c_b[split : 2 * split] if c_b is not None else None,
            )
            extracted["v"] = (
                c_w[2 * split :],
                c_b[2 * split :] if c_b is not None else None,
            )

        if hasattr(old_attn, "c_proj") and extracted["o"][0] is None:
            c = old_attn.c_proj
            w_o = c.weight.data.clone()
            if w_o.ndim == 2 and w_o.shape[0] < w_o.shape[1]:
                w_o = w_o.t().contiguous()
            b_o = getattr(c, "bias", None)
            extracted["o"] = (w_o, b_o.data.clone() if b_o is not None else None)

        # Nested BERT/RoBERTa
        if hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
            s = old_attn.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                if extracted[key][0] is None:
                    m = getattr(s, sub, None)
                    if m is not None and hasattr(m, "weight"):
                        b = getattr(m, "bias", None)
                        extracted[key] = (
                            m.weight.data.clone(),
                            b.data.clone() if b is not None else None,
                        )
            if (
                extracted["o"][0] is None
                and hasattr(old_attn, "output")
                and hasattr(old_attn.output, "dense")
            ):
                d = old_attn.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat modules
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        for key, names in named_map.items():
            if extracted[key][0] is None:
                for n in names:
                    if hasattr(old_attn, n):
                        m = getattr(old_attn, n)
                        if m is not None and hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # Apply SVD init
        for key, target in (
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ):
            w, b = extracted[key]
            if w is not None:
                try:
                    self._svd_init(
                        w,
                        b,
                        target,
                        whitening_L=self.whitening.get(key, None),
                    )
                except Exception as e:
                    logger.warning("Aegis-Hybrid SVD init failed for '%s': %s", key, e)
            elif key == "o" and self.is_nested_encoder:
                # Identity for nested encoder output (bypassed externally)
                with torch.no_grad():
                    r = min(target.out_features, target.in_features, target.rank)
                    target.down.weight.data.zero_()
                    target.up.weight.data.zero_()
                    target.down.weight.data[:r, :r].fill_diagonal_(1.0)
                    target.up.weight.data[:r, :r].fill_diagonal_(1.0)

        # storage accounting
        try:
            for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

    # ------------------------------------------------------------------
    # AH-FIX-06: comprehensive mask normalization
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(attention_mask, B, L, device) -> torch.Tensor | None:
        if attention_mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape == (B, L):
                return m
            if m.shape[0] == B and m.shape[1] != L:
                m = m[:, :L]
                if m.shape[1] == L:
                    return m
        elif attention_mask.dim() == 3:
            m = _to_bool(attention_mask)
            if m.shape[0] == B and m.shape[-1] == L:
                return m[:, -1, :]
        elif attention_mask.dim() == 4:
            m = _to_bool(attention_mask)
            m = m.any(dim=-2)
            if m.dim() == 3:
                m = m.squeeze(1)
            if m.shape == (B, L):
                return m
        return None

    # ------------------------------------------------------------------
    # AH-FIX-07: correct past_len from cache shape
    # ------------------------------------------------------------------
    @staticmethod
    def _read_past_kv(past_key_value, layer_idx=0):
        if past_key_value is None:
            return None, None
        if hasattr(past_key_value, "key_cache") and hasattr(past_key_value, "value_cache"):
            if not past_key_value.key_cache:
                return None, None
            idx = min(layer_idx, len(past_key_value.key_cache) - 1)
            return past_key_value.key_cache[idx], past_key_value.value_cache[idx]
        if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
            return past_key_value[0], past_key_value[1]
        return None, None

    # ------------------------------------------------------------------
    # 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,
    ):
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, N, _ = hidden_states.shape
        device = hidden_states.device
        Dh = self.head_dim

        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]

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

        # ---- RoPE ----
        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_k_probe, _ = self._read_past_kv(past_key_value, layer_idx)
            past_len_rope = 0 if past_k_probe is None else past_k_probe.shape[-2]
            cos = cos[..., past_len_rope : past_len_rope + N, :]
            sin = sin[..., past_len_rope : past_len_rope + N, :]
            Q, K_new = apply_rotary_pos_emb(Q, K_new, cos, sin)

        # ---- KV cache ----
        past_k, past_v = self._read_past_kv(past_key_value, layer_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)
            past_len = past_k.shape[-2]
        else:
            K_cat, V_cat = K_new, V_new
            past_len = 0

        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]

        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        # ---- mask ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

        # ---- local sliding window ----
        local_out, attn_local = _local_window_attention_chunked(
            Q,
            K_full,
            V_full,
            window=self.local_window,
            token_valid=token_valid,
            is_decoder=self.is_decoder and not is_cross,
            past_len=past_len,
            dropout_p=self.config.dropout,
            training=self.training,
        )

        # ---- global linear attention (optional) ----
        if self.use_global:
            global_out = _global_linear_attention(
                Q,
                K_full,
                V_full,
                self.kernel_map,
                token_valid,
            )
            beta = torch.sigmoid(self.fusion_beta).view(1, self.num_heads, 1, 1)
            # β → local, (1-β) → global
            fused = beta * local_out + (1.0 - beta) * global_out
        else:
            fused = local_out

        # ---- head_mask ----
        if head_mask is not None:
            hm = head_mask.to(fused.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(-1)
            fused = fused * hm

        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)

        if self.is_nested_encoder:
            final = fused
        else:
            final = self.out_proj(fused)

        final = torch.nan_to_num(final, nan=0.0, posinf=1e4, neginf=-1e4)

        return final.to(hidden_states.dtype), (attn_local if output_attentions else None)

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("q_proj", self.q_proj),
            ("k_proj", self.k_proj),
            ("v_proj", self.v_proj),
            ("out_proj", self.out_proj),
            ("kernel_map", self.kernel_map),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["fusion_beta"] = _sz(self.fusion_beta)
        out["total"] = sum(out.values())
        return out

Every 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.

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 storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("q_proj", self.q_proj),
        ("k_proj", self.k_proj),
        ("v_proj", self.v_proj),
        ("out_proj", self.out_proj),
        ("kernel_map", self.kernel_map),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["fusion_beta"] = _sz(self.fusion_beta)
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AegisHybridConfig (hidden_size: int = 768,
num_attention_heads: int = 12,
head_dim: int = 64,
num_key_value_heads: int | None = None,
local_window: int = 256,
rank: int = 128,
lora_rank: int = 32,
dropout: float = 0.1,
use_adaptive_pruning: bool = False,
gate_init_bias: float = 1.8,
is_decoder: bool = False,
use_global: bool = True,
kernel_feature_dim: int = 128,
kernel_seed: int = 1234,
chunk_size: int = 64)

Expand source code

@dataclass
class AegisHybridConfig:
    hidden_size: int = 768
    num_attention_heads: int = 12
    head_dim: int = 64
    num_key_value_heads: int | None = None
    local_window: int = 256
    rank: int = 128
    lora_rank: int = 32
    dropout: float = 0.1
    use_adaptive_pruning: bool = False
    gate_init_bias: float = 1.8
    is_decoder: bool = False
    use_global: bool = True
    kernel_feature_dim: int = 128
    kernel_seed: int = 1234
    chunk_size: int = 64

    def __post_init__(self) -> None:
        if self.num_key_value_heads is None:
            self.num_key_value_heads = self.num_attention_heads

AegisHybridConfig(hidden_size: 'int' = 768, num_attention_heads: 'int' = 12, head_dim: 'int' = 64, num_key_value_heads: 'int | None' = None, local_window: 'int' = 256, rank: 'int' = 128, lora_rank: 'int' = 32, dropout: 'float' = 0.1, use_adaptive_pruning: 'bool' = False, gate_init_bias: 'float' = 1.8, is_decoder: 'bool' = False, use_global: 'bool' = True, kernel_feature_dim: 'int' = 128, kernel_seed: 'int' = 1234, chunk_size: 'int' = 64)

Instance variables

var chunk_size : int

The type of the None singleton.

var dropout : float

The type of the None singleton.

var gate_init_bias : float

The type of the None singleton.

var head_dim : int

The type of the None singleton.

var hidden_size : int

The type of the None singleton.

var is_decoder : bool

The type of the None singleton.

var kernel_feature_dim : int

The type of the None singleton.

var kernel_seed : int

The type of the None singleton.

var local_window : int

The type of the None singleton.

var lora_rank : int

The type of the None singleton.

var num_attention_heads : int

The type of the None singleton.

var num_key_value_heads : int | None

The type of the None singleton.

var rank : int

The type of the None singleton.

var use_adaptive_pruning : bool

The type of the None singleton.

var use_global : bool

The type of the None singleton.

class AegisSpectronFP8Attention (config: Any | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("aegis_spectron_fp8", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("aegis_spectron_fp8_attn", supports_causal=True, supports_gqa=True)
class AegisSpectronFP8Attention(BaseAttention):
    is_custom_attention = True

    def __init__(
        self,
        config: Any | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ):
        if not isinstance(config, AegisSpectronFP8Config):
            cfg_src = config if config is not None else kwargs
            if isinstance(cfg_src, dict):
                h_size = hidden_size or cfg_src.get("hidden_size", cfg_src.get("dim", 768))
                n_heads = num_heads or cfg_src.get(
                    "num_attention_heads", cfg_src.get("num_heads", 12)
                )
                n_kv = cfg_src.get("num_key_value_heads", n_heads)
                h_dim = cfg_src.get("head_dim", h_size // max(1, n_heads))
            else:
                h_size = hidden_size or kwargs.get(
                    "hidden_size",
                    getattr(cfg_src, "hidden_size", getattr(cfg_src, "d_model", 768)),
                )
                n_heads = num_heads or kwargs.get(
                    "num_heads",
                    getattr(
                        cfg_src,
                        "num_attention_heads",
                        getattr(cfg_src, "num_heads", 12),
                    ),
                )
                n_kv = getattr(
                    cfg_src,
                    "num_key_value_heads",
                    getattr(cfg_src, "num_kv_heads", n_heads),
                )
                h_dim = getattr(
                    cfg_src,
                    "head_dim",
                    getattr(cfg_src, "kv_channels", h_size // max(1, n_heads)),
                )
            config = AegisSpectronFP8Config(
                hidden_size=h_size,
                num_heads=n_heads,
                num_key_value_heads=n_kv,
                head_dim=h_dim,
                window_size=getattr(cfg_src, "window_size", getattr(cfg_src, "local_window", 64)),
                svd_rank=getattr(cfg_src, "svd_rank", getattr(cfg_src, "rank", min(128, h_size))),
                lora_rank=getattr(cfg_src, "lora_rank", getattr(cfg_src, "lora_r", 32)),
                dropout=getattr(cfg_src, "dropout", 0.1),
                fp8_format=getattr(cfg_src, "fp8_format", "e4m3"),
                is_decoder=_detect_decoder(cfg_src, is_decoder, pretrained_attn),
            )
        super().__init__(config)
        self.config: AegisSpectronFP8Config = config

        self.hidden_size = config.hidden_size
        self.num_heads = config.num_heads
        self.head_dim = config.head_dim
        self.num_kv_heads = config.num_key_value_heads
        if self.num_heads % max(1, self.num_kv_heads) != 0:
            raise ValueError(
                f"num_heads ({self.num_heads}) must be divisible by "
                f"num_kv_heads ({self.num_kv_heads})."
            )
        self.kv_repeat = self.num_heads // max(1, self.num_kv_heads)
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim
        self.supports_gqa = True
        self.is_nested_encoder = False
        self.is_decoder = _detect_decoder(config, is_decoder, pretrained_attn)
        self.fp8_format = config.fp8_format
        self.rank = config.svd_rank
        self.lora_rank = config.lora_rank

        self.router = SpectronSpectralEntropyRouter(config)

        self.q_proj = FactorizedLinear(
            self.hidden_size,
            self.all_head_size,
            rank=self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.k_proj = FactorizedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            rank=self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.v_proj = FactorizedLinear(
            self.hidden_size,
            self.all_kv_head_size,
            rank=self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.out_proj = FactorizedLinear(
            self.all_head_size,
            self.hidden_size,
            rank=self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )

        # SP-FIX-15: separate Q/K low-rank projections.
        self.svd_U_q = nn.Parameter(torch.randn(self.all_head_size, self.rank) * 0.02)
        self.svd_U_k = nn.Parameter(torch.randn(self.all_head_size, self.rank) * 0.02)

        # Linear recurrence parameters (SP-FIX-06: no dead linear_A).
        self.linear_B = nn.Parameter(torch.ones(self.all_head_size))
        self.linear_C = nn.Parameter(torch.ones(self.all_head_size))
        self.decay_param = nn.Parameter(torch.full((self.all_head_size,), 4.5951))
        # SP-FIX-13: per-head gate.
        self.bidir_gate = nn.Parameter(torch.zeros(self.num_heads))
        self.linear_in = (
            nn.Linear(self.hidden_size, self.all_head_size, bias=False)
            if self.hidden_size != self.all_head_size
            else None
        )

        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 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

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

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

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

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

    # ------------------------------------------------------------------
    # SVD init of FactorizedLinear
    # ------------------------------------------------------------------
    @staticmethod
    def _svd_init(w: torch.Tensor, b: torch.Tensor | None, proj: FactorizedLinear) -> None:
        with torch.no_grad():
            W = w.detach().float()
            if W.ndim != 2:
                raise ValueError(f"Expected 2D weight, got shape {tuple(W.shape)}")
            U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            r = min(proj.rank, U.size(1), Vh.size(0))
            proj.lin1.weight.data.zero_()
            proj.lin2.weight.data.zero_()
            in_copy = min(W.shape[1], proj.lin1.weight.size(1))
            out_copy = min(W.shape[0], proj.lin2.weight.size(0))
            proj.lin1.weight.data[:r, :in_copy].copy_(Vh[:r, :in_copy].to(proj.lin1.weight.dtype))
            US = U[:out_copy, :r] @ torch.diag(S[:r])
            proj.lin2.weight.data[:out_copy, :r].copy_(US.to(proj.lin2.weight.dtype))
            if b is not None and proj.lin2.bias is not None:
                n = min(b.numel(), proj.lin2.bias.numel())
                proj.lin2.bias.data[:n].copy_(b.detach()[:n].to(proj.lin2.bias.dtype))
            elif proj.lin2.bias is not None:
                proj.lin2.bias.data.zero_()
            proj.lin1.weight.requires_grad = False
            proj.lin2.weight.requires_grad = False
            if proj.lin2.bias is not None:
                proj.lin2.bias.requires_grad = False
            if proj.lora_A is not None:
                nn.init.kaiming_uniform_(proj.lora_A, a=math.sqrt(5))
                nn.init.zeros_(proj.lora_B)

    # ------------------------------------------------------------------
    # Weight extraction (universal across architectures)
    # ------------------------------------------------------------------
    def _extract_proj_weights(
        self, old_attn: nn.Module
    ) -> dict[str, tuple[torch.Tensor | None, torch.Tensor | None]]:
        from neuralops.core.adapters import UniversalAdapter

        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(old_attn, "c_attn"):
            mod = old_attn.c_attn
            if hasattr(mod, "merge") and not getattr(mod, "merged", False):
                try:
                    mod.merge()
                except Exception:
                    pass
            base = getattr(mod, "base_layer", mod)
            c_w = getattr(base, "weight", getattr(mod, "weight", None))
            if c_w is not None:
                c_w = c_w.data
                is_conv1d = (
                    mod.__class__.__name__ == "Conv1D"
                    or base.__class__.__name__ == "Conv1D"
                    or hasattr(mod, "nf")
                    or hasattr(base, "nf")
                )
                if is_conv1d or (c_w.dim() == 2 and c_w.shape[1] == 3 * self.hidden_size):
                    c_w = c_w.t().contiguous()
                split = c_w.size(0) // 3
                bias_holder = getattr(base, "bias", getattr(mod, "bias", None))
                c_b = bias_holder.data if bias_holder is not None else None
                for i, p in enumerate(("q", "k", "v")):
                    extracted[p] = (
                        c_w[i * split : (i + 1) * split, :],
                        c_b[i * split : (i + 1) * split] if c_b is not None else None,
                    )

        # 2. GPT-2 c_proj
        if hasattr(old_attn, "c_proj"):
            mod = old_attn.c_proj
            if hasattr(mod, "merge") and not getattr(mod, "merged", False):
                try:
                    mod.merge()
                except Exception:
                    pass
            base = getattr(mod, "base_layer", mod)
            w = getattr(base, "weight", getattr(mod, "weight", None))
            if w is not None:
                w = w.data
                is_conv1d = (
                    mod.__class__.__name__ == "Conv1D"
                    or base.__class__.__name__ == "Conv1D"
                    or hasattr(mod, "nf")
                    or hasattr(base, "nf")
                )
                if is_conv1d or (
                    w.dim() == 2
                    and w.shape[0] != self.hidden_size
                    and w.shape[1] == self.hidden_size
                ):
                    w = w.t().contiguous()
                b = getattr(base, "bias", getattr(mod, "bias", None))
                extracted["o"] = (w, b.data if b is not None else None)

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

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

        # 5. Named projections
        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for n in names:
                    if hasattr(old_attn, n):
                        m = getattr(old_attn, n)
                        if hasattr(m, "weight"):
                            extracted[proj] = (m.weight.data, getattr(m, "bias", None))
                            break

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

        return extracted

    def copy_weights(self, old_attn: nn.Module) -> None:
        extracted = self._extract_proj_weights(old_attn)

        cls_name = old_attn.__class__.__name__
        has_out = any(
            hasattr(old_attn, a)
            for a in ("dense", "out_proj", "c_proj", "out_lin", "o_proj", "wo", "o")
        )
        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            (
                hasattr(old_attn, "self")
                and hasattr(old_attn.self, "query")
                and hasattr(old_attn, "output")
                and hasattr(old_attn.output, "dense")
            )
            or cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
            )
            or (cls_name.endswith("SelfAttention") and not has_out)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        missing = [p for p in ("q", "k", "v") if extracted[p][0] is None]
        if not self.is_nested_encoder and extracted["o"][0] is None:
            missing.append("o")
        if missing:
            raise RuntimeError(
                f"AegisSpectronFP8Attention.copy_weights: missing {missing} from "
                f"{cls_name}. Refusing random init."
            )

        for proj, target in (
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
        ):
            w, b = extracted[proj]
            if w is not None:
                self._svd_init(w, b, target)

        if self.is_nested_encoder:
            with torch.no_grad():
                self.out_proj.lin1.weight.data.zero_()
                self.out_proj.lin2.weight.data.zero_()
                r = min(self.out_proj.rank, self.hidden_size, self.all_head_size)
                self.out_proj.lin1.weight.data[:r, :r].fill_diagonal_(1.0)
                self.out_proj.lin2.weight.data[:r, :r].fill_diagonal_(1.0)
                if self.out_proj.lin2.bias is not None:
                    self.out_proj.lin2.bias.data.zero_()
        elif extracted["o"][0] is not None:
            self._svd_init(extracted["o"][0], extracted["o"][1], self.out_proj)

        # Init svd_U_q / svd_U_k from Q / K weights.
        with torch.no_grad():
            if extracted["q"][0] is not None:
                U, _, _ = torch.linalg.svd(extracted["q"][0].float(), full_matrices=False)
                r = min(self.rank, U.size(1), self.svd_U_q.size(1))
                self.svd_U_q.data[:, :r].copy_(U[: self.svd_U_q.size(0), :r].to(self.svd_U_q.dtype))
            if extracted["k"][0] is not None:
                U, _, _ = torch.linalg.svd(extracted["k"][0].float(), full_matrices=False)
                r = min(self.rank, U.size(1), self.svd_U_k.size(1))
                self.svd_U_k.data[:, :r].copy_(U[: self.svd_U_k.size(0), :r].to(self.svd_U_k.dtype))

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower():
            self.is_decoder = True

    # ------------------------------------------------------------------
    # Mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # Branch: local
    # ------------------------------------------------------------------
    def _local_branch(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        token_valid: torch.Tensor | None,
        past_len: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        return _local_window_attention(
            Q,
            K,
            V,
            window=self.config.window_size,
            token_valid=token_valid,
            is_decoder=self.is_decoder,
            past_len=past_len,
            dropout_p=self.config.dropout,
            training=self.training,
        )

    # ------------------------------------------------------------------
    # Branch: global (SP-FIX-17 — correct shapes)
    # ------------------------------------------------------------------
    def _global_branch(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        token_valid: torch.Tensor | None,
    ) -> torch.Tensor:
        """
        Returns context with the same shape as Q: (B, h, N, d).

        The low-rank projection collapses the head axis, so all linear
        attention math happens on (B, N, rank) / (B, N, h*d) tensors.
        Head structure is restored at the very end via a final view.
        """
        B, h, N, d = Q.shape
        D = h * d

        # Flatten heads → (B, N, h*d)
        Q_flat = Q.transpose(1, 2).contiguous().view(B, N, D)
        K_flat = K.transpose(1, 2).contiguous().view(B, N, D)
        V_flat = V.transpose(1, 2).contiguous().view(B, N, D)

        # Low-rank projection → (B, N, rank)
        Q_lr = torch.matmul(Q_flat, self.svd_U_q.to(Q_flat.dtype))
        K_lr = torch.matmul(K_flat, self.svd_U_k.to(K_flat.dtype))

        # Non-negative kernel feature map φ(x) = ELU(x) + 1
        phi_Q = F.elu(Q_lr) + 1.0  # (B, N, rank)
        phi_K = F.elu(K_lr) + 1.0

        # Run in fp32 for numerical stability (normalization can overflow in fp16).
        with torch.autocast(device_type=Q.device.type, enabled=False):
            phi_Q_f = phi_Q.float()
            phi_K_f = phi_K.float()
            V_f = V_flat.float()

            # SP-FIX-02: mask padding keys.
            if token_valid is not None:
                m = token_valid[:, :, None].to(phi_K_f.dtype)  # (B, N, 1)
                phi_K_f = phi_K_f * m
                V_f = V_f * m

            if self.is_decoder:
                out = _causal_global_linear_attention(
                    phi_Q_f,
                    phi_K_f,
                    V_f,
                    chunk_size=self.config.chunk_size,
                )
            else:
                out = _noncausal_global_linear_attention(phi_Q_f, phi_K_f, V_f)

            # (B, N, h*d) → (B, h, N, d)
            return (
                torch.nan_to_num(out, nan=0.0)
                .to(Q.dtype)
                .view(B, N, h, d)
                .transpose(1, 2)
                .contiguous()
            )

    # ------------------------------------------------------------------
    # Branch: linear recurrence (SP-FIX-05)
    # ------------------------------------------------------------------
    def _linear_branch(
        self,
        X: torch.Tensor,
        token_valid: torch.Tensor | None,
    ) -> torch.Tensor:
        if self.linear_in is not None:
            X = self.linear_in(X)
        B, N, D = X.shape

        if token_valid is not None:
            X = X * token_valid.unsqueeze(-1).to(X.dtype)

        u = X * self.linear_B.view(1, 1, D).to(X.dtype)
        gamma = torch.sigmoid(self.decay_param.to(X.dtype)).clamp(min=1e-4, max=1.0 - 1e-4)

        out_fwd = _state_recurrence_jit(u, gamma, reverse=False)

        if self.is_decoder:
            out_h = out_fwd.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        else:
            out_bwd = _state_recurrence_jit(u, gamma, reverse=True)
            gate = torch.sigmoid(self.bidir_gate.to(X.dtype)).view(1, self.num_heads, 1, 1)
            out_fwd_h = out_fwd.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
            out_bwd_h = out_bwd.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
            out_h = gate * out_fwd_h + (1.0 - gate) * out_bwd_h

        out_h = out_h * self.linear_C.view(1, self.num_heads, 1, self.head_dim).to(X.dtype)
        return out_h.contiguous()

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ):
        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
        token_valid_new = self._prepare_valid_mask(eff_mask, B, N_kv_new, device)

        # Routing on hidden states (self-attention only).
        alpha = self.router(hidden_states, token_valid=token_valid_new)
        alpha_loc = alpha[:, :, 0].view(B, 1, N, 1)
        alpha_glob = alpha[:, :, 1].view(B, 1, N, 1)
        alpha_lin = alpha[:, :, 2].view(B, 1, N, 1)

        # FP8 on projected activations (SP-FIX-01).
        fp8_fmt = self.fp8_format
        use_ste = self.training

        Q_new = fp8_cast_simulate(self.q_proj(hidden_states), format=fp8_fmt, use_ste=use_ste)
        K_new = fp8_cast_simulate(self.k_proj(kv_input), format=fp8_fmt, use_ste=use_ste)
        V_new = fp8_cast_simulate(self.v_proj(kv_input), format=fp8_fmt, use_ste=use_ste)

        Q = Q_new.view(B, N, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        K_new_h = (
            K_new.view(B, N_kv_new, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        )
        V_new_h = (
            V_new.view(B, N_kv_new, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        )

        # RoPE with past_len offset (SP-FIX-09).
        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_h = self._apply_rope(Q, K_new_h, cos, sin)

        # KV cache (SP-FIX-03).
        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(getattr(self, "layer_idx", 0), 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_h], dim=-2)
            V_cat = torch.cat([past_v, V_new_h], dim=-2)
        else:
            K_cat = K_new_h
            V_cat = V_new_h

        if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
            idx = getattr(self, "layer_idx", 0)
            past_key_value.update(K_new_h, V_new_h, idx)

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

        if past_len > 0 and token_valid_new is not None:
            past_valid = torch.ones(B, past_len, device=device, dtype=torch.bool)
            token_valid = torch.cat([past_valid, token_valid_new], dim=1)
        else:
            token_valid = token_valid_new

        K_full = repeat_kv(K_cat, self.kv_repeat) if self.kv_repeat > 1 else K_cat
        V_full = repeat_kv(V_cat, self.kv_repeat) if self.kv_repeat > 1 else V_cat

        out_loc, attn_loc = self._local_branch(Q, K_full, V_full, token_valid, past_len)
        out_glob = self._global_branch(Q, K_full, V_full, token_valid)
        out_lin = self._linear_branch(hidden_states, token_valid_new)
        target_dt = hidden_states.dtype
        fused = (
            alpha_loc.to(target_dt) * out_loc.to(target_dt)
            + alpha_glob.to(target_dt) * out_glob.to(target_dt)
            + alpha_lin.to(target_dt) * out_lin.to(target_dt)
        )

        # head_mask (SP-FIX-03).
        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)
            fused = fused * head_mask.to(fused.dtype)

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

        if self.is_nested_encoder:
            final = out
        else:
            final = self.out_proj(out)
        final = torch.nan_to_num(final, nan=0.0, posinf=1e4, neginf=-1e4)

        attn_ret = attn_loc if output_attentions else None
        return final.to(hidden_states.dtype), attn_ret

    # ------------------------------------------------------------------
    @staticmethod
    def _apply_rope(
        q: torch.Tensor,
        k: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        def _rotate_half(x: torch.Tensor) -> torch.Tensor:
            d = x.shape[-1]
            h = d // 2
            rot = torch.cat((-x[..., h : 2 * h], x[..., :h]), dim=-1)
            if d % 2 == 1:
                rot = torch.cat([rot, x[..., -1:]], dim=-1)
            return rot

        if cos.dim() == 3:
            cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
        while cos.dim() < q.dim():
            cos, sin = cos.unsqueeze(0), sin.unsqueeze(0)

        if cos.shape[-2] != q.shape[-2]:
            raise ValueError(f"RoPE seq len mismatch: cos {cos.shape[-2]} vs q {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:]
            q_rot = (q_rot * cos.to(q_rot.dtype)) + (_rotate_half(q_rot) * sin.to(q_rot.dtype))
            k_rot = (k_rot * cos.to(k_rot.dtype)) + (_rotate_half(k_rot) * sin.to(k_rot.dtype))
            q = torch.cat([q_rot, q_pass], dim=-1) if q_pass.numel() > 0 else q_rot
            k = torch.cat([k_rot, k_pass], dim=-1) if k_pass.numel() > 0 else k_rot
        else:
            cos_cut, sin_cut = cos[..., :head_dim], sin[..., :head_dim]
            q = (q * cos_cut.to(q.dtype)) + (_rotate_half(q) * sin_cut.to(q.dtype))
            k = (k * cos_cut.to(k.dtype)) + (_rotate_half(k) * sin_cut.to(k.dtype))
        return q, k

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.zeros((), device=self.out_proj.lin1.weight.device)

Every 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

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

Expand source code

@property
def k(self):
    return self.k_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

Expand source code

@property
def o(self):
    return self.out_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

Expand source code

@property
def q(self):
    return self.q_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

Expand source code

@property
def v(self):
    return self.v_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 get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.zeros((), device=self.out_proj.lin1.weight.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class AegisSpectronFP8Config (**kwargs: Any)

Expand source code

class AegisSpectronFP8Config:
    def __init__(self, **kwargs: Any):
        self.hidden_size = int(
            kwargs.get("hidden_size", kwargs.get("d_model", kwargs.get("dim", 768)))
        )
        self.num_heads = int(
            kwargs.get("num_heads", kwargs.get("num_attention_heads", kwargs.get("n_heads", 4)))
        )
        self.num_key_value_heads = int(
            kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads", self.num_heads))
        )
        self.head_dim = int(
            kwargs.get(
                "head_dim",
                kwargs.get("kv_channels", self.hidden_size // max(1, self.num_heads)),
            )
        )
        self.window_size = int(
            kwargs.get("window_size", kwargs.get("local_window", kwargs.get("block_size", 64)))
        )
        self.svd_rank = int(kwargs.get("svd_rank", kwargs.pop("rank", min(128, self.hidden_size))))
        self.lora_rank = int(kwargs.get("lora_rank", kwargs.get("lora_r", 32)))
        self.temperature = float(kwargs.get("temperature", 1.0))
        self.lambda_cost = float(kwargs.get("lambda_cost", 0.1))
        self.c_loc = float(kwargs.get("c_loc", 0.5))
        self.c_glob = float(kwargs.get("c_glob", 0.8))
        self.c_lin = float(kwargs.get("c_lin", 1.0))
        self.dropout = float(kwargs.get("dropout", 0.1))
        self.router_init_bias = tuple(kwargs.get("router_init_bias", (4.0, 0.0, 0.0)))
        self.is_decoder = bool(kwargs.get("is_decoder", False))
        self.fp8_format = str(kwargs.get("fp8_format", "e4m3"))
        self.chunk_size = int(kwargs.get("chunk_size", 64))

class Any (*args, **kwargs)

Expand source code

class Any(metaclass=_AnyMeta):
    """Special type indicating an unconstrained type.

    - Any is compatible with every type.
    - Any assumed to have all methods.
    - All values assumed to be instances of Any.

    Note that all the above statements are true from the point of view of
    static type checkers. At runtime, Any should not be used with instance
    checks.
    """

    def __new__(cls, *args, **kwargs):
        if cls is Any:
            raise TypeError("Any cannot be instantiated")
        return super().__new__(cls)

Special type indicating an unconstrained type.

• Any is compatible with every type.
• Any assumed to have all methods.
• All values assumed to be instances of Any.

Note that all the above statements are true from the point of view of static type checkers. At runtime, Any should not be used with instance checks.

class AutoConfig

Expand source code

class AutoConfig:
    r"""
    This is a generic configuration class that will be instantiated as one of the configuration classes of the library
    when created with the [`~AutoConfig.from_pretrained`] class method.

    This class cannot be instantiated directly using `__init__()` (throws an error).
    """

    def __init__(self) -> None:
        raise OSError(
            "AutoConfig is designed to be instantiated "
            "using the `AutoConfig.from_pretrained(pretrained_model_name_or_path)` method."
        )

    @classmethod
    def for_model(cls, model_type: str, *args, **kwargs) -> PreTrainedConfig:
        if model_type in CONFIG_MAPPING:
            config_class = CONFIG_MAPPING[model_type]
            return config_class(*args, **kwargs)
        raise ValueError(
            f"Unrecognized model identifier: {model_type}. Should contain one of {', '.join(CONFIG_MAPPING.keys())}"
        )

    @classmethod
    @replace_list_option_in_docstrings()
    def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], **kwargs):
        r"""
        Instantiate one of the configuration classes of the library from a pretrained model configuration.

        The configuration class to instantiate is selected based on the `model_type` property of the config object that
        is loaded, or when it's missing, by falling back to using pattern matching on `pretrained_model_name_or_path`:

        List options

        Args:
            pretrained_model_name_or_path (`str` or `os.PathLike`):
                Can be either:

                    - A string, the *model id* of a pretrained model configuration hosted inside a model repo on
                      huggingface.co.
                    - A path to a *directory* containing a configuration file saved using the
                      [`~PreTrainedConfig.save_pretrained`] method, or the [`~PreTrainedModel.save_pretrained`] method,
                      e.g., `./my_model_directory/`.
                    - a path to a saved configuration JSON *file*, e.g.,
                      `./my_model_directory/configuration.json`.
            cache_dir (`str` or `os.PathLike`, *optional*):
                Path to a directory in which a downloaded pretrained model configuration should be cached if the
                standard cache should not be used.
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force the (re-)download the model weights and configuration files and override the
                cached versions if they exist.
            proxies (`dict[str, str]`, *optional*):
                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
            revision (`str`, *optional*, defaults to `"main"`):
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
                identifier allowed by git.
            return_unused_kwargs (`bool`, *optional*, defaults to `False`):
                If `False`, then this function returns just the final configuration object.

                If `True`, then this functions returns a `Tuple(config, unused_kwargs)` where *unused_kwargs* is a
                dictionary consisting of the key/value pairs whose keys are not configuration attributes: i.e., the
                part of `kwargs` which has not been used to update `config` and is otherwise ignored.
            trust_remote_code (`bool`, *optional*, defaults to `False`):
                Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
                should only be set to `True` for repositories you trust and in which you have read the code, as it will
                execute code present on the Hub on your local machine.
            kwargs(additional keyword arguments, *optional*):
                The values in kwargs of any keys which are configuration attributes will be used to override the loaded
                values. Behavior concerning key/value pairs whose keys are *not* configuration attributes is controlled
                by the `return_unused_kwargs` keyword parameter.

        Examples:

        ```python
        >>> from transformers import AutoConfig

        >>> # Download configuration from huggingface.co and cache.
        >>> config = AutoConfig.from_pretrained("google-bert/bert-base-uncased")

        >>> # Download configuration from huggingface.co (user-uploaded) and cache.
        >>> config = AutoConfig.from_pretrained("dbmdz/bert-base-german-cased")

        >>> # If configuration file is in a directory (e.g., was saved using *save_pretrained('./test/saved_model/')*).
        >>> config = AutoConfig.from_pretrained("./test/bert_saved_model/")

        >>> # Load a specific configuration file.
        >>> config = AutoConfig.from_pretrained("./test/bert_saved_model/my_configuration.json")

        >>> # Change some config attributes when loading a pretrained config.
        >>> config = AutoConfig.from_pretrained("google-bert/bert-base-uncased", output_attentions=True, foo=False)
        >>> config.output_attentions
        True

        >>> config, unused_kwargs = AutoConfig.from_pretrained(
        ...     "google-bert/bert-base-uncased", output_attentions=True, foo=False, return_unused_kwargs=True
        ... )
        >>> config.output_attentions
        True

        >>> unused_kwargs
        {'foo': False}
        ```
        """
        kwargs["_from_auto"] = True
        kwargs["name_or_path"] = pretrained_model_name_or_path
        trust_remote_code = kwargs.pop("trust_remote_code", None)
        code_revision = kwargs.pop("code_revision", None)

        config_dict, unused_kwargs = PreTrainedConfig.get_config_dict(pretrained_model_name_or_path, **kwargs)
        has_remote_code = "auto_map" in config_dict and "AutoConfig" in config_dict["auto_map"]
        has_local_code = "model_type" in config_dict and config_dict["model_type"] in CONFIG_MAPPING
        explicit_local_code = has_local_code and not CONFIG_MAPPING[config_dict["model_type"]].__module__.startswith(
            "transformers."
        )
        if has_remote_code:
            class_ref = config_dict["auto_map"]["AutoConfig"]
            if "--" in class_ref:
                upstream_repo = class_ref.split("--")[0]
            else:
                upstream_repo = None
            trust_remote_code = resolve_trust_remote_code(
                trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code, upstream_repo
            )

        if has_remote_code and trust_remote_code and not explicit_local_code:
            config_class = get_class_from_dynamic_module(
                class_ref, pretrained_model_name_or_path, code_revision=code_revision, **kwargs
            )
            config_class.register_for_auto_class()
            return config_class.from_pretrained(pretrained_model_name_or_path, **kwargs)
        elif "model_type" in config_dict:
            # Apply heuristic: if model_type is mistral but layer_types is present, treat as ministral
            if config_dict["model_type"] == "mistral" and "layer_types" in config_dict:
                logger.info(
                    "Detected mistral model with layer_types, treating as ministral for alternating attention compatibility. "
                )
                config_dict["model_type"] = "ministral"

            try:
                config_class = CONFIG_MAPPING[config_dict["model_type"]]
            except KeyError:
                raise ValueError(
                    f"The checkpoint you are trying to load has model type `{config_dict['model_type']}` "
                    "but Transformers does not recognize this architecture. This could be because of an "
                    "issue with the checkpoint, or because your version of Transformers is out of date.\n\n"
                    "You can update Transformers with the command `pip install --upgrade transformers`. If this "
                    "does not work, and the checkpoint is very new, then there may not be a release version "
                    "that supports this model yet. In this case, you can get the most up-to-date code by installing "
                    "Transformers from source with the command "
                    "`pip install git+https://github.com/huggingface/transformers.git`"
                )
            return config_class.from_dict(config_dict, **unused_kwargs)

        raise ValueError(
            f"Unrecognized model in {pretrained_model_name_or_path}. "
            f"Should have a `model_type` key in its {CONFIG_NAME}."
        )

    @staticmethod
    def register(model_type, config, exist_ok=False) -> None:
        """
        Register a new configuration for this class.

        Args:
            model_type (`str`): The model type like "bert" or "gpt".
            config ([`PreTrainedConfig`]): The config to register.
        """
        if issubclass(config, PreTrainedConfig) and config.model_type != model_type:
            raise ValueError(
                "The config you are passing has a `model_type` attribute that is not consistent with the model type "
                f"you passed (config has {config.model_type} and you passed {model_type}. Fix one of those so they "
                "match!"
            )
        CONFIG_MAPPING.register(model_type, config, exist_ok=exist_ok)

This is a generic configuration class that will be instantiated as one of the configuration classes of the library when created with the [~AutoConfig.from_pretrained] class method.

This class cannot be instantiated directly using __init__() (throws an error).

Static methods

def for_model(model_type: str, *args, **kwargs) ‑> transformers.configuration_utils.PreTrainedConfig

def from_pretrained(pretrained_model_name_or_path: str | os.PathLike[str], **kwargs)

Instantiate one of the configuration classes of the library from a pretrained model configuration.

The configuration class to instantiate is selected based on the model_type property of the config object that is loaded, or when it's missing, by falling back to using pattern matching on pretrained_model_name_or_path:

- **EvollaModel** -- [`EvollaConfig`] (EvollaConfig model)
- **afmoe** -- [`AfmoeConfig`] (AfmoeConfig model)
- **aimv2** -- [`Aimv2Config`] (Aimv2Config model)
- **aimv2_text_model** -- [`Aimv2TextConfig`] (Aimv2TextConfig model)
- **aimv2_vision_model** -- [`Aimv2VisionConfig`] (Aimv2VisionConfig model)
- **albert** -- [`AlbertConfig`] (AlbertConfig model)
- **align** -- [`AlignConfig`] (AlignConfig model)
- **align_text_model** -- [`AlignTextConfig`] (AlignTextConfig model)
- **align_vision_model** -- [`AlignVisionConfig`] (AlignVisionConfig model)
- **altclip** -- [`AltCLIPConfig`] (AltCLIPConfig model)
- **altclip_text_model** -- [`AltCLIPTextConfig`] (AltCLIPTextConfig model)
- **altclip_vision_model** -- [`AltCLIPVisionConfig`] (AltCLIPVisionConfig model)
- **apertus** -- [`ApertusConfig`] (ApertusConfig model)
- **arcee** -- [`ArceeConfig`] (ArceeConfig model)
- **aria** -- [`AriaConfig`] (AriaConfig model)
- **aria_text** -- [`AriaTextConfig`] (AriaTextConfig model)
- **audio-spectrogram-transformer** -- [`ASTConfig`] (ASTConfig model)
- **audioflamingo3** -- [`AudioFlamingo3Config`] (AudioFlamingo3Config model)
- **audioflamingo3_encoder** -- [`AudioFlamingo3EncoderConfig`] (AudioFlamingo3EncoderConfig model)
- **autoformer** -- [`AutoformerConfig`] (AutoformerConfig model)
- **aya_vision** -- [`AyaVisionConfig`] (AyaVisionConfig model)
- **bamba** -- [`BambaConfig`] (BambaConfig model)
- **bark** -- [`BarkConfig`] (BarkConfig model)
- **bart** -- [`BartConfig`] (BartConfig model)
- **beit** -- [`BeitConfig`] (BeitConfig model)
- **bert** -- [`BertConfig`] (BertConfig model)
- **bert-generation** -- [`BertGenerationConfig`] (BertGenerationConfig model)
- **big_bird** -- [`BigBirdConfig`] (BigBirdConfig model)
- **bigbird_pegasus** -- [`BigBirdPegasusConfig`] (BigBirdPegasusConfig model)
- **biogpt** -- [`BioGptConfig`] (BioGptConfig model)
- **bit** -- [`BitConfig`] (BitConfig model)
- **bitnet** -- [`BitNetConfig`] (BitNetConfig model)
- **blenderbot** -- [`BlenderbotConfig`] (BlenderbotConfig model)
- **blenderbot-small** -- [`BlenderbotSmallConfig`] (BlenderbotSmallConfig model)
- **blip** -- [`BlipConfig`] (BlipConfig model)
- **blip-2** -- [`Blip2Config`] (Blip2Config model)
- **blip_2_qformer** -- [`Blip2QFormerConfig`] (Blip2QFormerConfig model)
- **blip_2_vision_model** -- [`Blip2VisionConfig`] (Blip2VisionConfig model)
- **blip_text_model** -- [`BlipTextConfig`] (BlipTextConfig model)
- **blip_vision_model** -- [`BlipVisionConfig`] (BlipVisionConfig model)
- **bloom** -- [`BloomConfig`] (BloomConfig model)
- **blt** -- [`BltConfig`] (BltConfig model)
- **blt_global_transformer** -- [`BltGlobalTransformerConfig`] (BltGlobalTransformerConfig model)
- **blt_local_decoder** -- [`BltLocalDecoderConfig`] (BltLocalDecoderConfig model)
- **blt_local_encoder** -- [`BltLocalEncoderConfig`] (BltLocalEncoderConfig model)
- **blt_patcher** -- [`BltPatcherConfig`] (BltPatcherConfig model)
- **bridgetower** -- [`BridgeTowerConfig`] (BridgeTowerConfig model)
- **bridgetower_text_model** -- [`BridgeTowerTextConfig`] (BridgeTowerTextConfig model)
- **bridgetower_vision_model** -- [`BridgeTowerVisionConfig`] (BridgeTowerVisionConfig model)
- **bros** -- [`BrosConfig`] (BrosConfig model)
- **camembert** -- [`CamembertConfig`] (CamembertConfig model)
- **canine** -- [`CanineConfig`] (CanineConfig model)
- **chameleon** -- [`ChameleonConfig`] (ChameleonConfig model)
- **chameleon_vqgan** -- [`ChameleonVQVAEConfig`] (ChameleonVQVAEConfig model)
- **chinese_clip** -- [`ChineseCLIPConfig`] (ChineseCLIPConfig model)
- **chinese_clip_text_model** -- [`ChineseCLIPTextConfig`] (ChineseCLIPTextConfig model)
- **chinese_clip_vision_model** -- [`ChineseCLIPVisionConfig`] (ChineseCLIPVisionConfig model)
- **chmv2** -- [`CHMv2Config`] (CHMv2Config model)
- **clap** -- [`ClapConfig`] (ClapConfig model)
- **clap_audio_model** -- [`ClapAudioConfig`] (ClapAudioConfig model)
- **clap_text_model** -- [`ClapTextConfig`] (ClapTextConfig model)
- **clip** -- [`CLIPConfig`] (CLIPConfig model)
- **clip_text_model** -- [`CLIPTextConfig`] (CLIPTextConfig model)
- **clip_vision_model** -- [`CLIPVisionConfig`] (CLIPVisionConfig model)
- **clipseg** -- [`CLIPSegConfig`] (CLIPSegConfig model)
- **clipseg_text_model** -- [`CLIPSegTextConfig`] (CLIPSegTextConfig model)
- **clipseg_vision_model** -- [`CLIPSegVisionConfig`] (CLIPSegVisionConfig model)
- **clvp** -- [`ClvpConfig`] (ClvpConfig model)
- **clvp_decoder** -- [`ClvpDecoderConfig`] (ClvpDecoderConfig model)
- **clvp_encoder** -- [`ClvpEncoderConfig`] (ClvpEncoderConfig model)
- **codegen** -- [`CodeGenConfig`] (CodeGenConfig model)
- **cohere** -- [`CohereConfig`] (CohereConfig model)
- **cohere2** -- [`Cohere2Config`] (Cohere2Config model)
- **cohere2_vision** -- [`Cohere2VisionConfig`] (Cohere2VisionConfig model)
- **cohere_asr** -- [`CohereAsrConfig`] (CohereAsrConfig model)
- **colmodernvbert** -- [`ColModernVBertConfig`] (ColModernVBertConfig model)
- **colpali** -- [`ColPaliConfig`] (ColPaliConfig model)
- **colqwen2** -- [`ColQwen2Config`] (ColQwen2Config model)
- **conditional_detr** -- [`ConditionalDetrConfig`] (ConditionalDetrConfig model)
- **convbert** -- [`ConvBertConfig`] (ConvBertConfig model)
- **convnext** -- [`ConvNextConfig`] (ConvNextConfig model)
- **convnextv2** -- [`ConvNextV2Config`] (ConvNextV2Config model)
- **cpmant** -- [`CpmAntConfig`] (CpmAntConfig model)
- **csm** -- [`CsmConfig`] (CsmConfig model)
- **csm_depth_decoder_model** -- [`CsmDepthDecoderConfig`] (CsmDepthDecoderConfig model)
- **ctrl** -- [`CTRLConfig`] (CTRLConfig model)
- **cvt** -- [`CvtConfig`] (CvtConfig model)
- **cwm** -- [`CwmConfig`] (CwmConfig model)
- **d_fine** -- [`DFineConfig`] (DFineConfig model)
- **dab-detr** -- [`DabDetrConfig`] (DabDetrConfig model)
- **dac** -- [`DacConfig`] (DacConfig model)
- **data2vec-audio** -- [`Data2VecAudioConfig`] (Data2VecAudioConfig model)
- **data2vec-text** -- [`Data2VecTextConfig`] (Data2VecTextConfig model)
- **data2vec-vision** -- [`Data2VecVisionConfig`] (Data2VecVisionConfig model)
- **dbrx** -- [`DbrxConfig`] (DbrxConfig model)
- **deberta** -- [`DebertaConfig`] (DebertaConfig model)
- **deberta-v2** -- [`DebertaV2Config`] (DebertaV2Config model)
- **decision_transformer** -- [`DecisionTransformerConfig`] (DecisionTransformerConfig model)
- **deepseek_v2** -- [`DeepseekV2Config`] (DeepseekV2Config model)
- **deepseek_v3** -- [`DeepseekV3Config`] (DeepseekV3Config model)
- **deepseek_vl** -- [`DeepseekVLConfig`] (DeepseekVLConfig model)
- **deepseek_vl_hybrid** -- [`DeepseekVLHybridConfig`] (DeepseekVLHybridConfig model)
- **deformable_detr** -- [`DeformableDetrConfig`] (DeformableDetrConfig model)
- **deimv2** -- [`Deimv2Config`] (Deimv2Config model)
- **deit** -- [`DeiTConfig`] (DeiTConfig model)
- **depth_anything** -- [`DepthAnythingConfig`] (DepthAnythingConfig model)
- **depth_pro** -- [`DepthProConfig`] (DepthProConfig model)
- **detr** -- [`DetrConfig`] (DetrConfig model)
- **dia** -- [`DiaConfig`] (DiaConfig model)
- **dia_decoder** -- [`DiaDecoderConfig`] (DiaDecoderConfig model)
- **dia_encoder** -- [`DiaEncoderConfig`] (DiaEncoderConfig model)
- **diffllama** -- [`DiffLlamaConfig`] (DiffLlamaConfig model)
- **dinat** -- [`DinatConfig`] (DinatConfig model)
- **dinov2** -- [`Dinov2Config`] (Dinov2Config model)
- **dinov2_with_registers** -- [`Dinov2WithRegistersConfig`] (Dinov2WithRegistersConfig model)
- **dinov3_convnext** -- [`DINOv3ConvNextConfig`] (DINOv3ConvNextConfig model)
- **dinov3_vit** -- [`DINOv3ViTConfig`] (DINOv3ViTConfig model)
- **distilbert** -- [`DistilBertConfig`] (DistilBertConfig model)
- **doge** -- [`DogeConfig`] (DogeConfig model)
- **donut-swin** -- [`DonutSwinConfig`] (DonutSwinConfig model)
- **dots1** -- [`Dots1Config`] (Dots1Config model)
- **dpr** -- [`DPRConfig`] (DPRConfig model)
- **dpt** -- [`DPTConfig`] (DPTConfig model)
- **edgetam** -- [`EdgeTamConfig`] (EdgeTamConfig model)
- **edgetam_video** -- [`EdgeTamVideoConfig`] (EdgeTamVideoConfig model)
- **edgetam_vision_model** -- [`EdgeTamVisionConfig`] (EdgeTamVisionConfig model)
- **efficientloftr** -- [`EfficientLoFTRConfig`] (EfficientLoFTRConfig model)
- **efficientnet** -- [`EfficientNetConfig`] (EfficientNetConfig model)
- **electra** -- [`ElectraConfig`] (ElectraConfig model)
- **emu3** -- [`Emu3Config`] (Emu3Config model)
- **emu3_text_model** -- [`Emu3TextConfig`] (Emu3TextConfig model)
- **emu3_vqgan** -- [`Emu3VQVAEConfig`] (Emu3VQVAEConfig model)
- **encodec** -- [`EncodecConfig`] (EncodecConfig model)
- **encoder-decoder** -- [`EncoderDecoderConfig`] (EncoderDecoderConfig model)
- **eomt** -- [`EomtConfig`] (EomtConfig model)
- **eomt_dinov3** -- [`EomtDinov3Config`] (EomtDinov3Config model)
- **ernie** -- [`ErnieConfig`] (ErnieConfig model)
- **ernie4_5** -- [`Ernie4_5Config`] (Ernie4_5Config model)
- **ernie4_5_moe** -- [`Ernie4_5_MoeConfig`] (Ernie4_5_MoeConfig model)
- **ernie4_5_vl_moe** -- [`Ernie4_5_VLMoeConfig`] (Ernie4_5_VLMoeConfig model)
- **ernie4_5_vl_moe_text** -- [`Ernie4_5_VLMoeTextConfig`] (Ernie4_5_VLMoeTextConfig model)
- **ernie4_5_vl_moe_vision** -- [`Ernie4_5_VLMoeVisionConfig`] (Ernie4_5_VLMoeVisionConfig model)
- **esm** -- [`EsmConfig`] (EsmConfig model)
- **eurobert** -- [`EuroBertConfig`] (EuroBertConfig model)
- **evolla** -- [`EvollaConfig`] (EvollaConfig model)
- **exaone4** -- [`Exaone4Config`] (Exaone4Config model)
- **exaone_moe** -- [`ExaoneMoeConfig`] (ExaoneMoeConfig model)
- **falcon** -- [`FalconConfig`] (FalconConfig model)
- **falcon_h1** -- [`FalconH1Config`] (FalconH1Config model)
- **falcon_mamba** -- [`FalconMambaConfig`] (FalconMambaConfig model)
- **fast_vlm** -- [`FastVlmConfig`] (FastVlmConfig model)
- **fastspeech2_conformer** -- [`FastSpeech2ConformerConfig`] (FastSpeech2ConformerConfig model)
- **fastspeech2_conformer_hifigan** -- [`FastSpeech2ConformerHifiGanConfig`] (FastSpeech2ConformerHifiGanConfig model)
- **fastspeech2_conformer_with_hifigan** -- [`FastSpeech2ConformerWithHifiGanConfig`] (FastSpeech2ConformerWithHifiGanConfig model)
- **flaubert** -- [`FlaubertConfig`] (FlaubertConfig model)
- **flava** -- [`FlavaConfig`] (FlavaConfig model)
- **flava_image_model** -- [`FlavaImageConfig`] (FlavaImageConfig model)
- **flava_multimodal_model** -- [`FlavaMultimodalConfig`] (FlavaMultimodalConfig model)
- **flava_text_model** -- [`FlavaTextConfig`] (FlavaTextConfig model)
- **flex_olmo** -- [`FlexOlmoConfig`] (FlexOlmoConfig model)
- **florence2** -- [`Florence2Config`] (Florence2Config model)
- **florence_vision** -- [`Florence2VisionConfig`] (Florence2VisionConfig model)
- **fnet** -- [`FNetConfig`] (FNetConfig model)
- **focalnet** -- [`FocalNetConfig`] (FocalNetConfig model)
- **fsmt** -- [`FSMTConfig`] (FSMTConfig model)
- **funnel** -- [`FunnelConfig`] (FunnelConfig model)
- **fuyu** -- [`FuyuConfig`] (FuyuConfig model)
- **gemma** -- [`GemmaConfig`] (GemmaConfig model)
- **gemma2** -- [`Gemma2Config`] (Gemma2Config model)
- **gemma3** -- [`Gemma3Config`] (Gemma3Config model)
- **gemma3_text** -- [`Gemma3TextConfig`] (Gemma3TextConfig model)
- **gemma3n** -- [`Gemma3nConfig`] (Gemma3nConfig model)
- **gemma3n_audio** -- [`Gemma3nAudioConfig`] (Gemma3nAudioConfig model)
- **gemma3n_text** -- [`Gemma3nTextConfig`] (Gemma3nTextConfig model)
- **gemma3n_vision** -- [`Gemma3nVisionConfig`] (Gemma3nVisionConfig model)
- **gemma4** -- [`Gemma4Config`] (Gemma4Config model)
- **gemma4_audio** -- [`Gemma4AudioConfig`] (Gemma4AudioConfig model)
- **gemma4_text** -- [`Gemma4TextConfig`] (Gemma4TextConfig model)
- **gemma4_vision** -- [`Gemma4VisionConfig`] (Gemma4VisionConfig model)
- **git** -- [`GitConfig`] (GitConfig model)
- **git_vision_model** -- [`GitVisionConfig`] (GitVisionConfig model)
- **glm** -- [`GlmConfig`] (GlmConfig model)
- **glm4** -- [`Glm4Config`] (Glm4Config model)
- **glm46v** -- [`Glm46VConfig`] (Glm46VConfig model)
- **glm4_moe** -- [`Glm4MoeConfig`] (Glm4MoeConfig model)
- **glm4_moe_lite** -- [`Glm4MoeLiteConfig`] (Glm4MoeLiteConfig model)
- **glm4v** -- [`Glm4vConfig`] (Glm4vConfig model)
- **glm4v_moe** -- [`Glm4vMoeConfig`] (Glm4vMoeConfig model)
- **glm4v_moe_text** -- [`Glm4vMoeTextConfig`] (Glm4vMoeTextConfig model)
- **glm4v_moe_vision** -- [`Glm4vMoeVisionConfig`] (Glm4vMoeVisionConfig model)
- **glm4v_text** -- [`Glm4vTextConfig`] (Glm4vTextConfig model)
- **glm4v_vision** -- [`Glm4vVisionConfig`] (Glm4vVisionConfig model)
- **glm_image** -- [`GlmImageConfig`] (GlmImageConfig model)
- **glm_image_text** -- [`GlmImageTextConfig`] (GlmImageTextConfig model)
- **glm_image_vision** -- [`GlmImageVisionConfig`] (GlmImageVisionConfig model)
- **glm_image_vqmodel** -- [`GlmImageVQVAEConfig`] (GlmImageVQVAEConfig model)
- **glm_moe_dsa** -- [`GlmMoeDsaConfig`] (GlmMoeDsaConfig model)
- **glm_ocr** -- [`GlmOcrConfig`] (GlmOcrConfig model)
- **glm_ocr_text** -- [`GlmOcrTextConfig`] (GlmOcrTextConfig model)
- **glm_ocr_vision** -- [`GlmOcrVisionConfig`] (GlmOcrVisionConfig model)
- **glmasr** -- [`GlmAsrConfig`] (GlmAsrConfig model)
- **glmasr_encoder** -- [`GlmAsrEncoderConfig`] (GlmAsrEncoderConfig model)
- **glpn** -- [`GLPNConfig`] (GLPNConfig model)
- **got_ocr2** -- [`GotOcr2Config`] (GotOcr2Config model)
- **gpt-sw3** -- [`GPT2Config`] (GPT2Config model)
- **gpt2** -- [`GPT2Config`] (GPT2Config model)
- **gpt_bigcode** -- [`GPTBigCodeConfig`] (GPTBigCodeConfig model)
- **gpt_neo** -- [`GPTNeoConfig`] (GPTNeoConfig model)
- **gpt_neox** -- [`GPTNeoXConfig`] (GPTNeoXConfig model)
- **gpt_neox_japanese** -- [`GPTNeoXJapaneseConfig`] (GPTNeoXJapaneseConfig model)
- **gpt_oss** -- [`GptOssConfig`] (GptOssConfig model)
- **gptj** -- [`GPTJConfig`] (GPTJConfig model)
- **granite** -- [`GraniteConfig`] (GraniteConfig model)
- **granite_speech** -- [`GraniteSpeechConfig`] (GraniteSpeechConfig model)
- **granite_speech_encoder** -- [`GraniteSpeechEncoderConfig`] (GraniteSpeechEncoderConfig model)
- **granitemoe** -- [`GraniteMoeConfig`] (GraniteMoeConfig model)
- **granitemoehybrid** -- [`GraniteMoeHybridConfig`] (GraniteMoeHybridConfig model)
- **granitemoeshared** -- [`GraniteMoeSharedConfig`] (GraniteMoeSharedConfig model)
- **grounding-dino** -- [`GroundingDinoConfig`] (GroundingDinoConfig model)
- **groupvit** -- [`GroupViTConfig`] (GroupViTConfig model)
- **groupvit_text_model** -- [`GroupViTTextConfig`] (GroupViTTextConfig model)
- **groupvit_vision_model** -- [`GroupViTVisionConfig`] (GroupViTVisionConfig model)
- **helium** -- [`HeliumConfig`] (HeliumConfig model)
- **hgnet_v2** -- [`HGNetV2Config`] (HGNetV2Config model)
- **hiera** -- [`HieraConfig`] (HieraConfig model)
- **higgs_audio_v2** -- [`HiggsAudioV2Config`] (HiggsAudioV2Config model)
- **higgs_audio_v2_tokenizer** -- [`HiggsAudioV2TokenizerConfig`] (HiggsAudioV2TokenizerConfig model)
- **hubert** -- [`HubertConfig`] (HubertConfig model)
- **hunyuan_v1_dense** -- [`HunYuanDenseV1Config`] (HunYuanDenseV1Config model)
- **hunyuan_v1_moe** -- [`HunYuanMoEV1Config`] (HunYuanMoEV1Config model)
- **hy_v3** -- [`HYV3Config`] (HYV3Config model)
- **ibert** -- [`IBertConfig`] (IBertConfig model)
- **idefics** -- [`IdeficsConfig`] (IdeficsConfig model)
- **idefics2** -- [`Idefics2Config`] (Idefics2Config model)
- **idefics2_perceiver** -- [`Idefics2PerceiverConfig`] (Idefics2PerceiverConfig model)
- **idefics2_vision** -- [`Idefics2VisionConfig`] (Idefics2VisionConfig model)
- **idefics3** -- [`Idefics3Config`] (Idefics3Config model)
- **idefics3_vision** -- [`Idefics3VisionConfig`] (Idefics3VisionConfig model)
- **idefics_perciever** -- [`IdeficsPerceiverConfig`] (IdeficsPerceiverConfig model)
- **idefics_vision** -- [`IdeficsVisionConfig`] (IdeficsVisionConfig model)
- **ijepa** -- [`IJepaConfig`] (IJepaConfig model)
- **imagegpt** -- [`ImageGPTConfig`] (ImageGPTConfig model)
- **informer** -- [`InformerConfig`] (InformerConfig model)
- **instructblip** -- [`InstructBlipConfig`] (InstructBlipConfig model)
- **instructblip_qformer** -- [`InstructBlipQFormerConfig`] (InstructBlipQFormerConfig model)
- **instructblip_vision_model** -- [`InstructBlipVisionConfig`] (InstructBlipVisionConfig model)
- **instructblipvideo** -- [`InstructBlipVideoConfig`] (InstructBlipVideoConfig model)
- **instructblipvideo_qformer** -- [`InstructBlipVideoQFormerConfig`] (InstructBlipVideoQFormerConfig model)
- **instructblipvideo_vision_model** -- [`InstructBlipVideoVisionConfig`] (InstructBlipVideoVisionConfig model)
- **internvl** -- [`InternVLConfig`] (InternVLConfig model)
- **internvl_vision** -- [`InternVLVisionConfig`] (InternVLVisionConfig model)
- **jais2** -- [`Jais2Config`] (Jais2Config model)
- **jamba** -- [`JambaConfig`] (JambaConfig model)
- **janus** -- [`JanusConfig`] (JanusConfig model)
- **janus_vision_model** -- [`JanusVisionConfig`] (JanusVisionConfig model)
- **janus_vqgan** -- [`JanusVQVAEConfig`] (JanusVQVAEConfig model)
- **jetmoe** -- [`JetMoeConfig`] (JetMoeConfig model)
- **jina_embeddings_v3** -- [`JinaEmbeddingsV3Config`] (JinaEmbeddingsV3Config model)
- **kosmos-2** -- [`Kosmos2Config`] (Kosmos2Config model)
- **kosmos-2.5** -- [`Kosmos2_5Config`] (Kosmos2_5Config model)
- **kosmos_2_5_text_model** -- [`Kosmos2_5TextConfig`] (Kosmos2_5TextConfig model)
- **kosmos_2_5_vision_model** -- [`Kosmos2_5VisionConfig`] (Kosmos2_5VisionConfig model)
- **kosmos_2_text_model** -- [`Kosmos2TextConfig`] (Kosmos2TextConfig model)
- **kosmos_2_vision_model** -- [`Kosmos2VisionConfig`] (Kosmos2VisionConfig model)
- **kyutai_speech_to_text** -- [`KyutaiSpeechToTextConfig`] (KyutaiSpeechToTextConfig model)
- **laguna** -- [`LagunaConfig`] (LagunaConfig model)
- **lasr_ctc** -- [`LasrCTCConfig`] (LasrCTCConfig model)
- **lasr_encoder** -- [`LasrEncoderConfig`] (LasrEncoderConfig model)
- **layoutlm** -- [`LayoutLMConfig`] (LayoutLMConfig model)
- **layoutlmv2** -- [`LayoutLMv2Config`] (LayoutLMv2Config model)
- **layoutlmv3** -- [`LayoutLMv3Config`] (LayoutLMv3Config model)
- **layoutxlm** -- [`LayoutXLMConfig`] (LayoutXLMConfig model)
- **led** -- [`LEDConfig`] (LEDConfig model)
- **levit** -- [`LevitConfig`] (LevitConfig model)
- **lfm2** -- [`Lfm2Config`] (Lfm2Config model)
- **lfm2_moe** -- [`Lfm2MoeConfig`] (Lfm2MoeConfig model)
- **lfm2_vl** -- [`Lfm2VlConfig`] (Lfm2VlConfig model)
- **lightglue** -- [`LightGlueConfig`] (LightGlueConfig model)
- **lighton_ocr** -- [`LightOnOcrConfig`] (LightOnOcrConfig model)
- **lilt** -- [`LiltConfig`] (LiltConfig model)
- **llama** -- [`LlamaConfig`] (LlamaConfig model)
- **llama4** -- [`Llama4Config`] (Llama4Config model)
- **llama4_text** -- [`Llama4TextConfig`] (Llama4TextConfig model)
- **llama4_vision_model** -- [`Llama4VisionConfig`] (Llama4VisionConfig model)
- **llava** -- [`LlavaConfig`] (LlavaConfig model)
- **llava_next** -- [`LlavaNextConfig`] (LlavaNextConfig model)
- **llava_next_video** -- [`LlavaNextVideoConfig`] (LlavaNextVideoConfig model)
- **llava_onevision** -- [`LlavaOnevisionConfig`] (LlavaOnevisionConfig model)
- **longcat_flash** -- [`LongcatFlashConfig`] (LongcatFlashConfig model)
- **longformer** -- [`LongformerConfig`] (LongformerConfig model)
- **longt5** -- [`LongT5Config`] (LongT5Config model)
- **luke** -- [`LukeConfig`] (LukeConfig model)
- **lw_detr** -- [`LwDetrConfig`] (LwDetrConfig model)
- **lw_detr_vit** -- [`LwDetrViTConfig`] (LwDetrViTConfig model)
- **lxmert** -- [`LxmertConfig`] (LxmertConfig model)
- **m2m_100** -- [`M2M100Config`] (M2M100Config model)
- **mamba** -- [`MambaConfig`] (MambaConfig model)
- **mamba2** -- [`Mamba2Config`] (Mamba2Config model)
- **marian** -- [`MarianConfig`] (MarianConfig model)
- **markuplm** -- [`MarkupLMConfig`] (MarkupLMConfig model)
- **mask2former** -- [`Mask2FormerConfig`] (Mask2FormerConfig model)
- **maskformer** -- [`MaskFormerConfig`] (MaskFormerConfig model)
- **maskformer-swin** -- [`MaskFormerSwinConfig`] (MaskFormerSwinConfig model)
- **mbart** -- [`MBartConfig`] (MBartConfig model)
- **megatron-bert** -- [`MegatronBertConfig`] (MegatronBertConfig model)
- **metaclip_2** -- [`MetaClip2Config`] (MetaClip2Config model)
- **metaclip_2_text_model** -- [`MetaClip2TextConfig`] (MetaClip2TextConfig model)
- **metaclip_2_vision_model** -- [`MetaClip2VisionConfig`] (MetaClip2VisionConfig model)
- **mgp-str** -- [`MgpstrConfig`] (MgpstrConfig model)
- **mimi** -- [`MimiConfig`] (MimiConfig model)
- **minicpmv4_6** -- [`MiniCPMV4_6Config`] (MiniCPMV4_6Config model)
- **minicpmv4_6_vision** -- [`MiniCPMV4_6VisionConfig`] (MiniCPMV4_6VisionConfig model)
- **minimax** -- [`MiniMaxConfig`] (MiniMaxConfig model)
- **minimax_m2** -- [`MiniMaxM2Config`] (MiniMaxM2Config model)
- **ministral** -- [`MinistralConfig`] (MinistralConfig model)
- **ministral3** -- [`Ministral3Config`] (Ministral3Config model)
- **mistral** -- [`MistralConfig`] (MistralConfig model)
- **mistral3** -- [`Mistral3Config`] (Mistral3Config model)
- **mistral4** -- [`Mistral4Config`] (Mistral4Config model)
- **mixtral** -- [`MixtralConfig`] (MixtralConfig model)
- **mlcd** -- [`MLCDVisionConfig`] (MLCDVisionConfig model)
- **mlcd_vision_model** -- [`MLCDVisionConfig`] (MLCDVisionConfig model)
- **mllama** -- [`MllamaConfig`] (MllamaConfig model)
- **mllama_text_model** -- [`MllamaTextConfig`] (MllamaTextConfig model)
- **mllama_vision_model** -- [`MllamaVisionConfig`] (MllamaVisionConfig model)
- **mm-grounding-dino** -- [`MMGroundingDinoConfig`] (MMGroundingDinoConfig model)
- **mobilebert** -- [`MobileBertConfig`] (MobileBertConfig model)
- **mobilenet_v1** -- [`MobileNetV1Config`] (MobileNetV1Config model)
- **mobilenet_v2** -- [`MobileNetV2Config`] (MobileNetV2Config model)
- **mobilevit** -- [`MobileViTConfig`] (MobileViTConfig model)
- **mobilevitv2** -- [`MobileViTV2Config`] (MobileViTV2Config model)
- **modernbert** -- [`ModernBertConfig`] (ModernBertConfig model)
- **modernbert-decoder** -- [`ModernBertDecoderConfig`] (ModernBertDecoderConfig model)
- **modernvbert** -- [`ModernVBertConfig`] (ModernVBertConfig model)
- **moonshine** -- [`MoonshineConfig`] (MoonshineConfig model)
- **moonshine_streaming** -- [`MoonshineStreamingConfig`] (MoonshineStreamingConfig model)
- **moonshine_streaming_encoder** -- [`MoonshineStreamingEncoderConfig`] (MoonshineStreamingEncoderConfig model)
- **moshi** -- [`MoshiConfig`] (MoshiConfig model)
- **moshi_depth** -- [`MoshiDepthConfig`] (MoshiDepthConfig model)
- **mpnet** -- [`MPNetConfig`] (MPNetConfig model)
- **mpt** -- [`MptConfig`] (MptConfig model)
- **mra** -- [`MraConfig`] (MraConfig model)
- **mt5** -- [`MT5Config`] (MT5Config model)
- **musicflamingo** -- [`MusicFlamingoConfig`] (MusicFlamingoConfig model)
- **musicgen** -- [`MusicgenConfig`] (MusicgenConfig model)
- **musicgen_decoder** -- [`MusicgenDecoderConfig`] (MusicgenDecoderConfig model)
- **musicgen_melody** -- [`MusicgenMelodyConfig`] (MusicgenMelodyConfig model)
- **musicgen_melody_decoder** -- [`MusicgenMelodyDecoderConfig`] (MusicgenMelodyDecoderConfig model)
- **mvp** -- [`MvpConfig`] (MvpConfig model)
- **nanochat** -- [`NanoChatConfig`] (NanoChatConfig model)
- **nemotron** -- [`NemotronConfig`] (NemotronConfig model)
- **nemotron_h** -- [`NemotronHConfig`] (NemotronHConfig model)
- **nllb-moe** -- [`NllbMoeConfig`] (NllbMoeConfig model)
- **nomic_bert** -- [`NomicBertConfig`] (NomicBertConfig model)
- **nougat** -- [`NougatConfig`] (NougatConfig model)
- **nystromformer** -- [`NystromformerConfig`] (NystromformerConfig model)
- **olmo** -- [`OlmoConfig`] (OlmoConfig model)
- **olmo2** -- [`Olmo2Config`] (Olmo2Config model)
- **olmo3** -- [`Olmo3Config`] (Olmo3Config model)
- **olmo_hybrid** -- [`OlmoHybridConfig`] (OlmoHybridConfig model)
- **olmoe** -- [`OlmoeConfig`] (OlmoeConfig model)
- **omdet-turbo** -- [`OmDetTurboConfig`] (OmDetTurboConfig model)
- **oneformer** -- [`OneFormerConfig`] (OneFormerConfig model)
- **openai-gpt** -- [`OpenAIGPTConfig`] (OpenAIGPTConfig model)
- **openai_privacy_filter** -- [`OpenAIPrivacyFilterConfig`] (OpenAIPrivacyFilterConfig model)
- **opt** -- [`OPTConfig`] (OPTConfig model)
- **ovis2** -- [`Ovis2Config`] (Ovis2Config model)
- **owlv2** -- [`Owlv2Config`] (Owlv2Config model)
- **owlv2_text_model** -- [`Owlv2TextConfig`] (Owlv2TextConfig model)
- **owlv2_vision_model** -- [`Owlv2VisionConfig`] (Owlv2VisionConfig model)
- **owlvit** -- [`OwlViTConfig`] (OwlViTConfig model)
- **owlvit_text_model** -- [`OwlViTTextConfig`] (OwlViTTextConfig model)
- **owlvit_vision_model** -- [`OwlViTVisionConfig`] (OwlViTVisionConfig model)
- **paddleocr_vl** -- [`PaddleOCRVLConfig`] (PaddleOCRVLConfig model)
- **paddleocr_vl_text** -- [`PaddleOCRTextConfig`] (PaddleOCRTextConfig model)
- **paddleocr_vl_vision** -- [`PaddleOCRVisionConfig`] (PaddleOCRVisionConfig model)
- **paligemma** -- [`PaliGemmaConfig`] (PaliGemmaConfig model)
- **parakeet_ctc** -- [`ParakeetCTCConfig`] (ParakeetCTCConfig model)
- **parakeet_encoder** -- [`ParakeetEncoderConfig`] (ParakeetEncoderConfig model)
- **patchtsmixer** -- [`PatchTSMixerConfig`] (PatchTSMixerConfig model)
- **patchtst** -- [`PatchTSTConfig`] (PatchTSTConfig model)
- **pe_audio** -- [`PeAudioConfig`] (PeAudioConfig model)
- **pe_audio_encoder** -- [`PeAudioEncoderConfig`] (PeAudioEncoderConfig model)
- **pe_audio_video** -- [`PeAudioVideoConfig`] (PeAudioVideoConfig model)
- **pe_audio_video_encoder** -- [`PeAudioVideoEncoderConfig`] (PeAudioVideoEncoderConfig model)
- **pe_video** -- [`PeVideoConfig`] (PeVideoConfig model)
- **pe_video_encoder** -- [`PeVideoEncoderConfig`] (PeVideoEncoderConfig model)
- **pegasus** -- [`PegasusConfig`] (PegasusConfig model)
- **pegasus_x** -- [`PegasusXConfig`] (PegasusXConfig model)
- **perceiver** -- [`PerceiverConfig`] (PerceiverConfig model)
- **perception_lm** -- [`PerceptionLMConfig`] (PerceptionLMConfig model)
- **persimmon** -- [`PersimmonConfig`] (PersimmonConfig model)
- **phi** -- [`PhiConfig`] (PhiConfig model)
- **phi3** -- [`Phi3Config`] (Phi3Config model)
- **phi4_multimodal** -- [`Phi4MultimodalConfig`] (Phi4MultimodalConfig model)
- **phi4_multimodal_audio** -- [`Phi4MultimodalAudioConfig`] (Phi4MultimodalAudioConfig model)
- **phi4_multimodal_vision** -- [`Phi4MultimodalVisionConfig`] (Phi4MultimodalVisionConfig model)
- **phimoe** -- [`PhimoeConfig`] (PhimoeConfig model)
- **pi0** -- [`PI0Config`] (PI0Config model)
- **pix2struct** -- [`Pix2StructConfig`] (Pix2StructConfig model)
- **pix2struct_text_model** -- [`Pix2StructTextConfig`] (Pix2StructTextConfig model)
- **pix2struct_vision_model** -- [`Pix2StructVisionConfig`] (Pix2StructVisionConfig model)
- **pixio** -- [`PixioConfig`] (PixioConfig model)
- **pixtral** -- [`PixtralVisionConfig`] (PixtralVisionConfig model)
- **plbart** -- [`PLBartConfig`] (PLBartConfig model)
- **poolformer** -- [`PoolFormerConfig`] (PoolFormerConfig model)
- **pop2piano** -- [`Pop2PianoConfig`] (Pop2PianoConfig model)
- **pp_chart2table** -- [`PPChart2TableConfig`] (PPChart2TableConfig model)
- **pp_doclayout_v2** -- [`PPDocLayoutV2Config`] (PPDocLayoutV2Config model)
- **pp_doclayout_v3** -- [`PPDocLayoutV3Config`] (PPDocLayoutV3Config model)
- **pp_lcnet** -- [`PPLCNetConfig`] (PPLCNetConfig model)
- **pp_lcnet_v3** -- [`PPLCNetV3Config`] (PPLCNetV3Config model)
- **pp_ocrv5_mobile_det** -- [`PPOCRV5MobileDetConfig`] (PPOCRV5MobileDetConfig model)
- **pp_ocrv5_mobile_rec** -- [`PPOCRV5MobileRecConfig`] (PPOCRV5MobileRecConfig model)
- **pp_ocrv5_server_det** -- [`PPOCRV5ServerDetConfig`] (PPOCRV5ServerDetConfig model)
- **pp_ocrv5_server_rec** -- [`PPOCRV5ServerRecConfig`] (PPOCRV5ServerRecConfig model)
- **prompt_depth_anything** -- [`PromptDepthAnythingConfig`] (PromptDepthAnythingConfig model)
- **prophetnet** -- [`ProphetNetConfig`] (ProphetNetConfig model)
- **pvt** -- [`PvtConfig`] (PvtConfig model)
- **pvt_v2** -- [`PvtV2Config`] (PvtV2Config model)
- **qianfan_ocr** -- [`QianfanOCRConfig`] (QianfanOCRConfig model)
- **qianfan_ocr_vision** -- [`QianfanOCRVisionConfig`] (QianfanOCRVisionConfig model)
- **qwen2** -- [`Qwen2Config`] (Qwen2Config model)
- **qwen2_5_omni** -- [`Qwen2_5OmniConfig`] (Qwen2_5OmniConfig model)
- **qwen2_5_omni_audio_encoder** -- [`Qwen2_5OmniAudioEncoderConfig`] (Qwen2_5OmniAudioEncoderConfig model)
- **qwen2_5_omni_bigvgan** -- [`Qwen2_5OmniBigVGANConfig`] (Qwen2_5OmniBigVGANConfig model)
- **qwen2_5_omni_dit** -- [`Qwen2_5OmniDiTConfig`] (Qwen2_5OmniDiTConfig model)
- **qwen2_5_omni_talker** -- [`Qwen2_5OmniTalkerConfig`] (Qwen2_5OmniTalkerConfig model)
- **qwen2_5_omni_text** -- [`Qwen2_5OmniTextConfig`] (Qwen2_5OmniTextConfig model)
- **qwen2_5_omni_thinker** -- [`Qwen2_5OmniThinkerConfig`] (Qwen2_5OmniThinkerConfig model)
- **qwen2_5_omni_token2wav** -- [`Qwen2_5OmniToken2WavConfig`] (Qwen2_5OmniToken2WavConfig model)
- **qwen2_5_omni_vision_encoder** -- [`Qwen2_5OmniVisionEncoderConfig`] (Qwen2_5OmniVisionEncoderConfig model)
- **qwen2_5_vl** -- [`Qwen2_5_VLConfig`] (Qwen2_5_VLConfig model)
- **qwen2_5_vl_text** -- [`Qwen2_5_VLTextConfig`] (Qwen2_5_VLTextConfig model)
- **qwen2_5_vl_vision** -- [`Qwen2_5_VLVisionConfig`] (Qwen2_5_VLVisionConfig model)
- **qwen2_audio** -- [`Qwen2AudioConfig`] (Qwen2AudioConfig model)
- **qwen2_audio_encoder** -- [`Qwen2AudioEncoderConfig`] (Qwen2AudioEncoderConfig model)
- **qwen2_moe** -- [`Qwen2MoeConfig`] (Qwen2MoeConfig model)
- **qwen2_vl** -- [`Qwen2VLConfig`] (Qwen2VLConfig model)
- **qwen2_vl_text** -- [`Qwen2VLTextConfig`] (Qwen2VLTextConfig model)
- **qwen2_vl_vision** -- [`Qwen2VLVisionConfig`] (Qwen2VLVisionConfig model)
- **qwen3** -- [`Qwen3Config`] (Qwen3Config model)
- **qwen3_5** -- [`Qwen3_5Config`] (Qwen3_5Config model)
- **qwen3_5_moe** -- [`Qwen3_5MoeConfig`] (Qwen3_5MoeConfig model)
- **qwen3_5_moe_text** -- [`Qwen3_5MoeTextConfig`] (Qwen3_5MoeTextConfig model)
- **qwen3_5_moe_vision** -- [`Qwen3_5MoeVisionConfig`] (Qwen3_5MoeVisionConfig model)
- **qwen3_5_text** -- [`Qwen3_5TextConfig`] (Qwen3_5TextConfig model)
- **qwen3_5_vision** -- [`Qwen3_5VisionConfig`] (Qwen3_5VisionConfig model)
- **qwen3_moe** -- [`Qwen3MoeConfig`] (Qwen3MoeConfig model)
- **qwen3_next** -- [`Qwen3NextConfig`] (Qwen3NextConfig model)
- **qwen3_omni_moe** -- [`Qwen3OmniMoeConfig`] (Qwen3OmniMoeConfig model)
- **qwen3_omni_moe_audio_encoder** -- [`Qwen3OmniMoeAudioEncoderConfig`] (Qwen3OmniMoeAudioEncoderConfig model)
- **qwen3_omni_moe_talker_code_predictor** -- [`Qwen3OmniMoeTalkerCodePredictorConfig`] (Qwen3OmniMoeTalkerCodePredictorConfig model)
- **qwen3_omni_moe_talker_text** -- [`Qwen3OmniMoeTalkerTextConfig`] (Qwen3OmniMoeTalkerTextConfig model)
- **qwen3_omni_moe_text** -- [`Qwen3OmniMoeTextConfig`] (Qwen3OmniMoeTextConfig model)
- **qwen3_omni_moe_thinker** -- [`Qwen3OmniMoeThinkerConfig`] (Qwen3OmniMoeThinkerConfig model)
- **qwen3_omni_moe_vision_encoder** -- [`Qwen3OmniMoeVisionEncoderConfig`] (Qwen3OmniMoeVisionEncoderConfig model)
- **qwen3_vl** -- [`Qwen3VLConfig`] (Qwen3VLConfig model)
- **qwen3_vl_moe** -- [`Qwen3VLMoeConfig`] (Qwen3VLMoeConfig model)
- **qwen3_vl_moe_text** -- [`Qwen3VLMoeTextConfig`] (Qwen3VLMoeTextConfig model)
- **qwen3_vl_moe_vision** -- [`Qwen3VLMoeVisionConfig`] (Qwen3VLMoeVisionConfig model)
- **qwen3_vl_text** -- [`Qwen3VLTextConfig`] (Qwen3VLTextConfig model)
- **qwen3_vl_vision** -- [`Qwen3VLVisionConfig`] (Qwen3VLVisionConfig model)
- **rag** -- [`RagConfig`] (RagConfig model)
- **recurrent_gemma** -- [`RecurrentGemmaConfig`] (RecurrentGemmaConfig model)
- **reformer** -- [`ReformerConfig`] (ReformerConfig model)
- **regnet** -- [`RegNetConfig`] (RegNetConfig model)
- **rembert** -- [`RemBertConfig`] (RemBertConfig model)
- **resnet** -- [`ResNetConfig`] (ResNetConfig model)
- **roberta** -- [`RobertaConfig`] (RobertaConfig model)
- **roberta-prelayernorm** -- [`RobertaPreLayerNormConfig`] (RobertaPreLayerNormConfig model)
- **roc_bert** -- [`RoCBertConfig`] (RoCBertConfig model)
- **roformer** -- [`RoFormerConfig`] (RoFormerConfig model)
- **rt_detr** -- [`RTDetrConfig`] (RTDetrConfig model)
- **rt_detr_resnet** -- [`RTDetrResNetConfig`] (RTDetrResNetConfig model)
- **rt_detr_v2** -- [`RTDetrV2Config`] (RTDetrV2Config model)
- **rwkv** -- [`RwkvConfig`] (RwkvConfig model)
- **sam** -- [`SamConfig`] (SamConfig model)
- **sam2** -- [`Sam2Config`] (Sam2Config model)
- **sam2_hiera_det_model** -- [`Sam2HieraDetConfig`] (Sam2HieraDetConfig model)
- **sam2_video** -- [`Sam2VideoConfig`] (Sam2VideoConfig model)
- **sam2_vision_model** -- [`Sam2VisionConfig`] (Sam2VisionConfig model)
- **sam3** -- [`Sam3Config`] (Sam3Config model)
- **sam3_detr_decoder** -- [`Sam3DETRDecoderConfig`] (Sam3DETRDecoderConfig model)
- **sam3_detr_encoder** -- [`Sam3DETREncoderConfig`] (Sam3DETREncoderConfig model)
- **sam3_geometry_encoder** -- [`Sam3GeometryEncoderConfig`] (Sam3GeometryEncoderConfig model)
- **sam3_lite_text** -- [`Sam3LiteTextConfig`] (Sam3LiteTextConfig model)
- **sam3_lite_text_detr_decoder** -- [`Sam3LiteTextDETRDecoderConfig`] (Sam3LiteTextDETRDecoderConfig model)
- **sam3_lite_text_detr_encoder** -- [`Sam3LiteTextDETREncoderConfig`] (Sam3LiteTextDETREncoderConfig model)
- **sam3_lite_text_geometry_encoder** -- [`Sam3LiteTextGeometryEncoderConfig`] (Sam3LiteTextGeometryEncoderConfig model)
- **sam3_lite_text_mask_decoder** -- [`Sam3LiteTextMaskDecoderConfig`] (Sam3LiteTextMaskDecoderConfig model)
- **sam3_lite_text_text_model** -- [`Sam3LiteTextTextConfig`] (Sam3LiteTextTextConfig model)
- **sam3_mask_decoder** -- [`Sam3MaskDecoderConfig`] (Sam3MaskDecoderConfig model)
- **sam3_tracker** -- [`Sam3TrackerConfig`] (Sam3TrackerConfig model)
- **sam3_tracker_video** -- [`Sam3TrackerVideoConfig`] (Sam3TrackerVideoConfig model)
- **sam3_video** -- [`Sam3VideoConfig`] (Sam3VideoConfig model)
- **sam3_vision_model** -- [`Sam3VisionConfig`] (Sam3VisionConfig model)
- **sam3_vit_model** -- [`Sam3ViTConfig`] (Sam3ViTConfig model)
- **sam_hq** -- [`SamHQConfig`] (SamHQConfig model)
- **sam_hq_vision_model** -- [`SamHQVisionConfig`] (SamHQVisionConfig model)
- **sam_vision_model** -- [`SamVisionConfig`] (SamVisionConfig model)
- **seamless_m4t** -- [`SeamlessM4TConfig`] (SeamlessM4TConfig model)
- **seamless_m4t_v2** -- [`SeamlessM4Tv2Config`] (SeamlessM4Tv2Config model)
- **seed_oss** -- [`SeedOssConfig`] (SeedOssConfig model)
- **segformer** -- [`SegformerConfig`] (SegformerConfig model)
- **seggpt** -- [`SegGptConfig`] (SegGptConfig model)
- **sew** -- [`SEWConfig`] (SEWConfig model)
- **sew-d** -- [`SEWDConfig`] (SEWDConfig model)
- **shieldgemma2** -- [`ShieldGemma2Config`] (ShieldGemma2Config model)
- **siglip** -- [`SiglipConfig`] (SiglipConfig model)
- **siglip2** -- [`Siglip2Config`] (Siglip2Config model)
- **siglip2_text_model** -- [`Siglip2TextConfig`] (Siglip2TextConfig model)
- **siglip2_vision_model** -- [`Siglip2VisionConfig`] (Siglip2VisionConfig model)
- **siglip_text_model** -- [`SiglipTextConfig`] (SiglipTextConfig model)
- **siglip_vision_model** -- [`SiglipVisionConfig`] (SiglipVisionConfig model)
- **slanet** -- [`SLANetConfig`] (SLANetConfig model)
- **slanext** -- [`SLANeXtConfig`] (SLANeXtConfig model)
- **smollm3** -- [`SmolLM3Config`] (SmolLM3Config model)
- **smolvlm** -- [`SmolVLMConfig`] (SmolVLMConfig model)
- **smolvlm_vision** -- [`SmolVLMVisionConfig`] (SmolVLMVisionConfig model)
- **solar_open** -- [`SolarOpenConfig`] (SolarOpenConfig model)
- **speech-encoder-decoder** -- [`SpeechEncoderDecoderConfig`] (SpeechEncoderDecoderConfig model)
- **speech_to_text** -- [`Speech2TextConfig`] (Speech2TextConfig model)
- **speecht5** -- [`SpeechT5Config`] (SpeechT5Config model)
- **speecht5_hifigan** -- [`SpeechT5HifiGanConfig`] (SpeechT5HifiGanConfig model)
- **splinter** -- [`SplinterConfig`] (SplinterConfig model)
- **squeezebert** -- [`SqueezeBertConfig`] (SqueezeBertConfig model)
- **stablelm** -- [`StableLmConfig`] (StableLmConfig model)
- **starcoder2** -- [`Starcoder2Config`] (Starcoder2Config model)
- **superglue** -- [`SuperGlueConfig`] (SuperGlueConfig model)
- **superpoint** -- [`SuperPointConfig`] (SuperPointConfig model)
- **swiftformer** -- [`SwiftFormerConfig`] (SwiftFormerConfig model)
- **swin** -- [`SwinConfig`] (SwinConfig model)
- **swin2sr** -- [`Swin2SRConfig`] (Swin2SRConfig model)
- **swinv2** -- [`Swinv2Config`] (Swinv2Config model)
- **switch_transformers** -- [`SwitchTransformersConfig`] (SwitchTransformersConfig model)
- **t5** -- [`T5Config`] (T5Config model)
- **t5_gemma_module** -- [`T5GemmaModuleConfig`] (T5GemmaModuleConfig model)
- **t5gemma** -- [`T5GemmaConfig`] (T5GemmaConfig model)
- **t5gemma2** -- [`T5Gemma2Config`] (T5Gemma2Config model)
- **t5gemma2_decoder** -- [`T5Gemma2DecoderConfig`] (T5Gemma2DecoderConfig model)
- **t5gemma2_encoder** -- [`T5Gemma2EncoderConfig`] (T5Gemma2EncoderConfig model)
- **t5gemma2_text** -- [`T5Gemma2TextConfig`] (T5Gemma2TextConfig model)
- **table-transformer** -- [`TableTransformerConfig`] (TableTransformerConfig model)
- **tapas** -- [`TapasConfig`] (TapasConfig model)
- **textnet** -- [`TextNetConfig`] (TextNetConfig model)
- **time_series_transformer** -- [`TimeSeriesTransformerConfig`] (TimeSeriesTransformerConfig model)
- **timesfm** -- [`TimesFmConfig`] (TimesFmConfig model)
- **timesfm2_5** -- [`TimesFm2_5Config`] (TimesFm2_5Config model)
- **timesformer** -- [`TimesformerConfig`] (TimesformerConfig model)
- **timm_backbone** -- [`TimmBackboneConfig`] (TimmBackboneConfig model)
- **timm_wrapper** -- [`TimmWrapperConfig`] (TimmWrapperConfig model)
- **trocr** -- [`TrOCRConfig`] (TrOCRConfig model)
- **tvp** -- [`TvpConfig`] (TvpConfig model)
- **udop** -- [`UdopConfig`] (UdopConfig model)
- **umt5** -- [`UMT5Config`] (UMT5Config model)
- **unispeech** -- [`UniSpeechConfig`] (UniSpeechConfig model)
- **unispeech-sat** -- [`UniSpeechSatConfig`] (UniSpeechSatConfig model)
- **univnet** -- [`UnivNetConfig`] (UnivNetConfig model)
- **upernet** -- [`UperNetConfig`] (UperNetConfig model)
- **uvdoc** -- [`UVDocConfig`] (UVDocConfig model)
- **uvdoc_backbone** -- [`UVDocBackboneConfig`] (UVDocBackboneConfig model)
- **vaultgemma** -- [`VaultGemmaConfig`] (VaultGemmaConfig model)
- **vibevoice_acoustic_tokenizer** -- [`VibeVoiceAcousticTokenizerConfig`] (VibeVoiceAcousticTokenizerConfig model)
- **vibevoice_acoustic_tokenizer_decoder** -- [`VibeVoiceAcousticTokenizerDecoderConfig`] (VibeVoiceAcousticTokenizerDecoderConfig model)
- **vibevoice_acoustic_tokenizer_encoder** -- [`VibeVoiceAcousticTokenizerEncoderConfig`] (VibeVoiceAcousticTokenizerEncoderConfig model)
- **vibevoice_asr** -- [`VibeVoiceAsrConfig`] (VibeVoiceAsrConfig model)
- **video_llama_3** -- [`VideoLlama3Config`] (VideoLlama3Config model)
- **video_llama_3_vision** -- [`VideoLlama3VisionConfig`] (VideoLlama3VisionConfig model)
- **video_llava** -- [`VideoLlavaConfig`] (VideoLlavaConfig model)
- **videomae** -- [`VideoMAEConfig`] (VideoMAEConfig model)
- **videomt** -- [`VideomtConfig`] (VideomtConfig model)
- **vilt** -- [`ViltConfig`] (ViltConfig model)
- **vipllava** -- [`VipLlavaConfig`] (VipLlavaConfig model)
- **vision-encoder-decoder** -- [`VisionEncoderDecoderConfig`] (VisionEncoderDecoderConfig model)
- **vision-text-dual-encoder** -- [`VisionTextDualEncoderConfig`] (VisionTextDualEncoderConfig model)
- **visual_bert** -- [`VisualBertConfig`] (VisualBertConfig model)
- **vit** -- [`ViTConfig`] (ViTConfig model)
- **vit_mae** -- [`ViTMAEConfig`] (ViTMAEConfig model)
- **vit_msn** -- [`ViTMSNConfig`] (ViTMSNConfig model)
- **vitdet** -- [`VitDetConfig`] (VitDetConfig model)
- **vitmatte** -- [`VitMatteConfig`] (VitMatteConfig model)
- **vitpose** -- [`VitPoseConfig`] (VitPoseConfig model)
- **vitpose_backbone** -- [`VitPoseBackboneConfig`] (VitPoseBackboneConfig model)
- **vits** -- [`VitsConfig`] (VitsConfig model)
- **vivit** -- [`VivitConfig`] (VivitConfig model)
- **vjepa2** -- [`VJEPA2Config`] (VJEPA2Config model)
- **voxtral** -- [`VoxtralConfig`] (VoxtralConfig model)
- **voxtral_encoder** -- [`VoxtralEncoderConfig`] (VoxtralEncoderConfig model)
- **voxtral_realtime** -- [`VoxtralRealtimeConfig`] (VoxtralRealtimeConfig model)
- **voxtral_realtime_encoder** -- [`VoxtralRealtimeEncoderConfig`] (VoxtralRealtimeEncoderConfig model)
- **voxtral_realtime_text** -- [`VoxtralRealtimeTextConfig`] (VoxtralRealtimeTextConfig model)
- **wav2vec2** -- [`Wav2Vec2Config`] (Wav2Vec2Config model)
- **wav2vec2-bert** -- [`Wav2Vec2BertConfig`] (Wav2Vec2BertConfig model)
- **wav2vec2-conformer** -- [`Wav2Vec2ConformerConfig`] (Wav2Vec2ConformerConfig model)
- **wavlm** -- [`WavLMConfig`] (WavLMConfig model)
- **whisper** -- [`WhisperConfig`] (WhisperConfig model)
- **xclip** -- [`XCLIPConfig`] (XCLIPConfig model)
- **xclip_text_model** -- [`XCLIPTextConfig`] (XCLIPTextConfig model)
- **xclip_vision_model** -- [`XCLIPVisionConfig`] (XCLIPVisionConfig model)
- **xcodec** -- [`XcodecConfig`] (XcodecConfig model)
- **xglm** -- [`XGLMConfig`] (XGLMConfig model)
- **xlm** -- [`XLMConfig`] (XLMConfig model)
- **xlm-roberta** -- [`XLMRobertaConfig`] (XLMRobertaConfig model)
- **xlm-roberta-xl** -- [`XLMRobertaXLConfig`] (XLMRobertaXLConfig model)
- **xlnet** -- [`XLNetConfig`] (XLNetConfig model)
- **xlstm** -- [`xLSTMConfig`] (xLSTMConfig model)
- **xmod** -- [`XmodConfig`] (XmodConfig model)
- **yolos** -- [`YolosConfig`] (YolosConfig model)
- **yoso** -- [`YosoConfig`] (YosoConfig model)
- **youtu** -- [`YoutuConfig`] (YoutuConfig model)
- **zamba** -- [`ZambaConfig`] (ZambaConfig model)
- **zamba2** -- [`Zamba2Config`] (Zamba2Config model)
- **zoedepth** -- [`ZoeDepthConfig`] (ZoeDepthConfig model)

Args -----= pretrained_model_name_or_path (str or os.PathLike): Can be either:

    - A string, the *model id* of a pretrained model configuration hosted inside a model repo on
      huggingface.co.
    - A path to a *directory* containing a configuration file saved using the
      [`~PreTrainedConfig.save_pretrained`] method, or the [`~PreTrainedModel.save_pretrained`] method,
      e.g., `./my_model_directory/`.
    - a path to a saved configuration JSON *file*, e.g.,
      `./my_model_directory/configuration.json`.

cache_dir (str or os.PathLike, optional): Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. force_download (bool, optional, defaults to False): Whether or not to force the (re-)download the model weights and configuration files and override the cached versions if they exist. proxies (dict[str, str], optional): A dictionary of proxy servers to use by protocol or endpoint, e.g., {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}. The proxies are used on each request. revision (str, optional, defaults to "main"): The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a git-based system for storing models and other artifacts on huggingface.co, so revision can be any identifier allowed by git. return_unused_kwargs (bool, optional, defaults to False): If False, then this function returns just the final configuration object.

If <code>True</code>, then this functions returns a <code>Tuple(config, unused\_kwargs)</code> where *unused_kwargs* is a
dictionary consisting of the key/value pairs whose keys are not configuration attributes: i.e., the
part of <code>kwargs</code> which has not been used to update <code>config</code> and is otherwise ignored.

trust_remote_code (bool, optional, defaults to False): Whether or not to allow for custom models defined on the Hub in their own modeling files. This option should only be set to True for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine. kwargs(additional keyword arguments, optional): The values in kwargs of any keys which are configuration attributes will be used to override the loaded values. Behavior concerning key/value pairs whose keys are not configuration attributes is controlled by the return_unused_kwargs keyword parameter. Examples:

>>> from transformers import AutoConfig

>>> # Download configuration from huggingface.co and cache.
>>> config = AutoConfig.from_pretrained("google-bert/bert-base-uncased")

>>> # Download configuration from huggingface.co (user-uploaded) and cache.
>>> config = AutoConfig.from_pretrained("dbmdz/bert-base-german-cased")

>>> # If configuration file is in a directory (e.g., was saved using *save_pretrained('./test/saved_model/')*).
>>> config = AutoConfig.from_pretrained("./test/bert_saved_model/")

>>> # Load a specific configuration file.
>>> config = AutoConfig.from_pretrained("./test/bert_saved_model/my_configuration.json")

>>> # Change some config attributes when loading a pretrained config.
>>> config = AutoConfig.from_pretrained("google-bert/bert-base-uncased", output_attentions=True, foo=False)
>>> config.output_attentions
True

>>> config, unused_kwargs = AutoConfig.from_pretrained(
...     "google-bert/bert-base-uncased", output_attentions=True, foo=False, return_unused_kwargs=True
... )
>>> config.output_attentions
True

>>> unused_kwargs
{'foo': False}

def register(model_type, config, exist_ok=False) ‑> None

Expand source code

@staticmethod
def register(model_type, config, exist_ok=False) -> None:
    """
    Register a new configuration for this class.

    Args:
        model_type (`str`): The model type like "bert" or "gpt".
        config ([`PreTrainedConfig`]): The config to register.
    """
    if issubclass(config, PreTrainedConfig) and config.model_type != model_type:
        raise ValueError(
            "The config you are passing has a `model_type` attribute that is not consistent with the model type "
            f"you passed (config has {config.model_type} and you passed {model_type}. Fix one of those so they "
            "match!"
        )
    CONFIG_MAPPING.register(model_type, config, exist_ok=exist_ok)

Register a new configuration for this class.

Args -----= model_type (str): The model type like "bert" or "gpt". config ([PreTrainedConfig]): The config to register.

class AutoModelForSequenceClassification (*args, **kwargs)

Expand source code

class AutoModelForSequenceClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING

This is a generic model class that will be instantiated as one of the model classes of the library (with a sequence classification head) when created with the [~AutoModelForSequenceClassification.from_pretrained] class method or the [~AutoModelForSequenceClassification.from_config] class method.

This class cannot be instantiated directly using __init__() (throws an error).

Ancestors

• transformers.models.auto.auto_factory._BaseAutoModelClass

Static methods

def from_config(config, **kwargs)

def from_pretrained(pretrained_model_name_or_path: str | os.PathLike[str], *model_args, **kwargs)

class AutoTokenizer

Expand source code

class AutoTokenizer:
    r"""
    This is a generic tokenizer class that will be instantiated as one of the tokenizer classes of the library when
    created with the [`AutoTokenizer.from_pretrained`] class method.

    This class cannot be instantiated directly using `__init__()` (throws an error).
    """

    def __init__(self):
        raise OSError(
            "AutoTokenizer is designed to be instantiated "
            "using the `AutoTokenizer.from_pretrained(pretrained_model_name_or_path)` method."
        )

    @classmethod
    @replace_list_option_in_docstrings(TOKENIZER_MAPPING_NAMES)
    def from_pretrained(
        cls, pretrained_model_name_or_path, *inputs, **kwargs
    ) -> TokenizersBackend | SentencePieceBackend:
        r"""
        Instantiate one of the tokenizer classes of the library from a pretrained model vocabulary.

        The tokenizer class to instantiate is selected based on the `model_type` property of the config object (either
        passed as an argument or loaded from `pretrained_model_name_or_path` if possible), or when it's missing, by
        falling back to using pattern matching on `pretrained_model_name_or_path`:

        List options

        Params:
            pretrained_model_name_or_path (`str` or `os.PathLike`):
                Can be either:

                    - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co.
                    - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved
                      using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
                    - a path to a single saved vocabulary file if and only if the tokenizer only requires a
                      single vocabulary file (like Bert or XLNet), e.g.: `./my_model_directory/vocab.txt`. (Not
                      applicable to all derived classes)
            inputs (additional positional arguments, *optional*):
                Will be passed along to the Tokenizer `__init__()` method.
            config ([`PreTrainedConfig`], *optional*)
                The configuration object used to determine the tokenizer class to instantiate.
            cache_dir (`str` or `os.PathLike`, *optional*):
                Path to a directory in which a downloaded pretrained model configuration should be cached if the
                standard cache should not be used.
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force the (re-)download the model weights and configuration files and override the
                cached versions if they exist.
            proxies (`dict[str, str]`, *optional*):
                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
            revision (`str`, *optional*, defaults to `"main"`):
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
                identifier allowed by git.
            subfolder (`str`, *optional*):
                In case the relevant files are located inside a subfolder of the model repo on huggingface.co (e.g. for
                facebook/rag-token-base), specify it here.
            tokenizer_type (`str`, *optional*):
                Tokenizer type to be loaded.
            backend (`str`, *optional*, defaults to `"tokenizers"`):
                Backend to use for tokenization. Valid options are:
                - `"tokenizers"`: Use the HuggingFace tokenizers library backend (default)
                - `"sentencepiece"`: Use the SentencePiece backend
            trust_remote_code (`bool`, *optional*, defaults to `False`):
                Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
                should only be set to `True` for repositories you trust and in which you have read the code, as it will
                execute code present on the Hub on your local machine.
            kwargs (additional keyword arguments, *optional*):
                Will be passed to the Tokenizer `__init__()` method. Can be used to set special tokens like
                `bos_token`, `eos_token`, `unk_token`, `sep_token`, `pad_token`, `cls_token`, `mask_token`,
                `additional_special_tokens`. See parameters in the `__init__()` for more details.

        Examples:

        ```python
        >>> from transformers import AutoTokenizer

        >>> # Download vocabulary from huggingface.co and cache.
        >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")

        >>> # Download vocabulary from huggingface.co (user-uploaded) and cache.
        >>> tokenizer = AutoTokenizer.from_pretrained("dbmdz/bert-base-german-cased")

        >>> # If vocabulary files are in a directory (e.g. tokenizer was saved using *save_pretrained('./test/saved_model/')*)
        >>> # tokenizer = AutoTokenizer.from_pretrained("./test/bert_saved_model/")

        >>> # Download vocabulary from huggingface.co and define model-specific arguments
        >>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/roberta-base", add_prefix_space=True)

        >>> # Explicitly use the tokenizers backend
        >>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer", backend="tokenizers")

        >>> # Explicitly use the sentencepiece backend
        >>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer", backend="sentencepiece")
        ```"""
        config = kwargs.pop("config", None)
        kwargs["_from_auto"] = True

        # V5: Always use fast tokenizers, ignore use_fast parameter
        _ = kwargs.pop("use_fast", None)
        tokenizer_type = kwargs.pop("tokenizer_type", None)
        trust_remote_code = kwargs.pop("trust_remote_code", None)
        gguf_file = kwargs.get("gguf_file")

        # First, let's see whether the tokenizer_type is passed so that we can leverage it
        if tokenizer_type is not None:
            tokenizer_class_name = TOKENIZER_MAPPING_NAMES.get(tokenizer_type, None)

            if tokenizer_class_name is None:
                raise ValueError(
                    f"Passed `tokenizer_type` {tokenizer_type} does not exist. `tokenizer_type` should be one of "
                    f"{', '.join(c for c in TOKENIZER_MAPPING_NAMES)}."
                )

            tokenizer_class = tokenizer_class_from_name(tokenizer_class_name)

            if tokenizer_class is None:
                raise ValueError(f"Tokenizer class {tokenizer_class_name} is not currently imported.")

            return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

        if gguf_file:
            gguf_path = cached_file(pretrained_model_name_or_path, gguf_file, **kwargs)
            config_dict = load_gguf_checkpoint(gguf_path, return_tensors=False)["config"]
            config = AutoConfig.for_model(**config_dict)
        elif config is None:
            try:
                config = AutoConfig.from_pretrained(
                    pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
                )
            except (ValueError, OSError):
                config = PreTrainedConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)

        config_model_type = config.model_type

        # Next, let's try to use the tokenizer_config file to get the tokenizer class.
        tokenizer_config = get_tokenizer_config(pretrained_model_name_or_path, **kwargs)
        tokenizer_config_class = tokenizer_config.get("tokenizer_class", None)

        # Check for auto_map early to handle dynamic tokenizers properly
        tokenizer_auto_map = None
        if "auto_map" in tokenizer_config:
            if isinstance(tokenizer_config["auto_map"], (tuple, list)):
                # Legacy format for dynamic tokenizers
                tokenizer_auto_map = tokenizer_config["auto_map"]
            else:
                tokenizer_auto_map = tokenizer_config["auto_map"].get("AutoTokenizer", None)

        # if there is a config, we can check that the tokenizer class != than model class.
        # Use the config class if it's a specialized tokenizer, otherwise fall back to TokenizersBackend.
        if (
            tokenizer_auto_map is None
            and tokenizer_config_class is not None
            and config_model_type is not None
            and config_model_type != ""
            and TOKENIZER_MAPPING_NAMES.get(config_model_type) is not None
            and (TOKENIZER_MAPPING_NAMES.get(config_model_type).removesuffix("Fast"))
            != (tokenizer_config_class.removesuffix("Fast"))
        ):
            registered_class_name = TOKENIZER_MAPPING_NAMES.get(config_model_type).removesuffix("Fast")
            if registered_class_name not in ("TokenizersBackend", "PythonBackend", "PreTrainedTokenizerFast"):
                # The auto-mapping has a real class but the Hub specifies a different specialized class so trust the Hub's class.
                tokenizer_class = tokenizer_class_from_name(tokenizer_config_class)
                if tokenizer_class is not None and tokenizer_class.__name__ not in (
                    "TokenizersBackend",
                    "PythonBackend",
                    "PreTrainedTokenizerFast",
                ):
                    return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

            if TokenizersBackend is not None:
                return TokenizersBackend.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

            raise ValueError(
                f"Tokenizer class '{tokenizer_config_class}' specified in the tokenizer config was not found. "
                f"The tokenizer may need to be converted or re-saved."
            )

        if "_commit_hash" in tokenizer_config:
            kwargs["_commit_hash"] = tokenizer_config["_commit_hash"]

        if tokenizer_config_class and tokenizer_config_class.endswith("Fast"):
            tokenizer_config_class = tokenizer_config_class[:-4]

        has_remote_code = tokenizer_auto_map is not None
        has_local_code = type(config) in TOKENIZER_MAPPING or (
            tokenizer_config_class is not None
            and (
                tokenizer_class_from_name(tokenizer_config_class) is not None
                or tokenizer_class_from_name(tokenizer_config_class + "Fast") is not None
            )
        )
        explicit_local_code = (
            has_local_code
            and type(config) not in TOKENIZER_MAPPING
            and (
                tokenizer_config_class is not None
                and not (
                    tokenizer_class_from_name(tokenizer_config_class)
                    or tokenizer_class_from_name(tokenizer_config_class + "Fast")
                ).__module__.startswith("transformers.")
            )
        )
        # V5: Skip remote tokenizer for custom models with incorrect hub tokenizer class
        if has_remote_code and config_model_type in MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS:
            has_remote_code = False
            tokenizer_auto_map = None

        if has_remote_code:
            # V5: Always prefer fast tokenizer (index 1), fallback to slow (index 0)
            if tokenizer_auto_map[1] is not None:
                class_ref = tokenizer_auto_map[1]
            else:
                class_ref = tokenizer_auto_map[0]
            if "--" in class_ref:
                upstream_repo = class_ref.split("--")[0]
            else:
                upstream_repo = None
            trust_remote_code = resolve_trust_remote_code(
                trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code, upstream_repo
            )

        if has_remote_code and trust_remote_code and not explicit_local_code:
            # BC v5: register *Fast aliases before remote code loads.
            if tokenizer_config_class:
                tokenizer_class_from_name(tokenizer_config_class.removesuffix("Fast"))
            tokenizer_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
            _ = kwargs.pop("code_revision", None)
            tokenizer_class.register_for_auto_class()
            return tokenizer_class.from_pretrained(
                pretrained_model_name_or_path, *inputs, trust_remote_code=trust_remote_code, **kwargs
            )
        elif tokenizer_config_class is not None:
            tokenizer_class_candidate = tokenizer_config_class
            tokenizer_class = tokenizer_class_from_name(tokenizer_class_candidate)
            if tokenizer_class is None and not tokenizer_class_candidate.endswith("Fast"):
                tokenizer_class = tokenizer_class_from_name(tokenizer_class_candidate + "Fast")
            if tokenizer_class is not None and tokenizer_class.__name__ == "PythonBackend":
                tokenizer_class = TokenizersBackend
            # Fallback to TokenizersBackend if the class wasn't found
            if tokenizer_class is None:
                tokenizer_class = TokenizersBackend

            return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
        elif getattr(config, "tokenizer_class", None):
            _class = config.tokenizer_class
            if "PreTrainedTokenizerFast" not in _class and _class.endswith("Fast"):
                _class = _class[:-4]
            tokenizer_class = tokenizer_class_from_name(_class)
            return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

        # Otherwise we have to be creative.
        # if model is an encoder decoder, the encoder tokenizer class is used by default
        if isinstance(config, EncoderDecoderConfig):
            if type(config.decoder) is not type(config.encoder):
                logger.warning(
                    f"The encoder model config class: {config.encoder.__class__} is different from the decoder model "
                    f"config class: {config.decoder.__class__}. It is not recommended to use the "
                    "`AutoTokenizer.from_pretrained()` method in this case. Please use the encoder and decoder "
                    "specific tokenizer classes."
                )
            config = config.encoder

        model_type = config_class_to_model_type(type(config).__name__) or getattr(config, "model_type", None)
        if model_type is not None:
            tokenizer_class = TOKENIZER_MAPPING.get(type(config), TokenizersBackend)
            if tokenizer_class is not None:
                return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

        # Fallback: try tokenizer_class from tokenizer_config.json
        tokenizer_config_class = tokenizer_config.get("tokenizer_class", None)
        if tokenizer_config_class is not None:
            if tokenizer_config_class != "TokenizersBackend" and tokenizer_config_class.endswith("Fast"):
                tokenizer_config_class = tokenizer_config_class[:-4]
            tokenizer_class = tokenizer_class_from_name(tokenizer_config_class)
            if tokenizer_class is None and not tokenizer_config_class.endswith("Fast"):
                tokenizer_class = tokenizer_class_from_name(tokenizer_config_class + "Fast")
            if tokenizer_class is not None and tokenizer_class.__name__ == "PythonBackend":
                tokenizer_class = TokenizersBackend
            if tokenizer_class is None:
                tokenizer_class = TokenizersBackend
            return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

        raise ValueError(
            f"Unrecognized configuration class {config.__class__} to build an AutoTokenizer.\n"
            f"Model type should be one of {', '.join(c.__name__ for c in TOKENIZER_MAPPING)}."
        )

    @staticmethod
    def register(
        config_class, tokenizer_class=None, slow_tokenizer_class=None, fast_tokenizer_class=None, exist_ok=False
    ):
        """
        Register a new tokenizer in this mapping.

        Args:
            config_class ([`PreTrainedConfig`]):
                The configuration corresponding to the model to register.
            tokenizer_class: The tokenizer class to register (V5 - preferred parameter).
            slow_tokenizer_class: (Deprecated) The slow tokenizer to register.
            fast_tokenizer_class: (Deprecated) The fast tokenizer to register.
        """
        if tokenizer_class is None:
            # Legacy: prefer fast over slow
            if fast_tokenizer_class is not None:
                tokenizer_class = fast_tokenizer_class
            elif slow_tokenizer_class is not None:
                tokenizer_class = slow_tokenizer_class
            else:
                raise ValueError("You need to pass a `tokenizer_class`")

        for candidate in (slow_tokenizer_class, fast_tokenizer_class, tokenizer_class):
            if candidate is not None:
                REGISTERED_TOKENIZER_CLASSES[candidate.__name__] = candidate

        if slow_tokenizer_class is not None and fast_tokenizer_class is not None:
            REGISTERED_FAST_ALIASES[slow_tokenizer_class.__name__] = fast_tokenizer_class

        TOKENIZER_MAPPING.register(config_class, tokenizer_class, exist_ok=exist_ok)

This is a generic tokenizer class that will be instantiated as one of the tokenizer classes of the library when created with the [AutoTokenizer.from_pretrained] class method.

This class cannot be instantiated directly using __init__() (throws an error).

Static methods

def from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs) ‑> transformers.tokenization_utils_tokenizers.TokenizersBackend | None

Instantiate one of the tokenizer classes of the library from a pretrained model vocabulary.

The tokenizer class to instantiate is selected based on the model_type property of the config object (either passed as an argument or loaded from pretrained_model_name_or_path if possible), or when it's missing, by falling back to using pattern matching on pretrained_model_name_or_path:

- **aimv2** -- [`CLIPTokenizer`] (Aimv2Config model)
- **albert** -- [`AlbertTokenizer`] (AlbertConfig model)
- **align** -- [`BertTokenizer`] (AlignConfig model)
- **audioflamingo3** -- [`Qwen2Tokenizer`] (AudioFlamingo3Config model)
- **aya_vision** -- [`CohereTokenizer`] (AyaVisionConfig model)
- **bark** -- [`BertTokenizer`] (BarkConfig model)
- **bart** -- [`RobertaTokenizer`] (BartConfig model)
- **bert** -- [`BertTokenizer`] (BertConfig model)
- **bert-generation** -- [`None`] (BertGenerationConfig model)
- **big_bird** -- [`BigBirdTokenizer`] (BigBirdConfig model)
- **bigbird_pegasus** -- [`PegasusTokenizer`] (BigBirdPegasusConfig model)
- **biogpt** -- [`BioGptTokenizer`] (BioGptConfig model)
- **blenderbot** -- [`BlenderbotTokenizer`] (BlenderbotConfig model)
- **blenderbot-small** -- [`BlenderbotSmallTokenizer`] (BlenderbotSmallConfig model)
- **blip** -- [`BertTokenizer`] (BlipConfig model)
- **blip-2** -- [`GPT2Tokenizer`] (Blip2Config model)
- **bridgetower** -- [`RobertaTokenizer`] (BridgeTowerConfig model)
- **bros** -- [`BertTokenizer`] (BrosConfig model)
- **camembert** -- [`CamembertTokenizer`] (CamembertConfig model)
- **canine** -- [`CanineTokenizer`] (CanineConfig model)
- **chameleon** -- [`TokenizersBackend`] (ChameleonConfig model)
- **chinese_clip** -- [`BertTokenizer`] (ChineseCLIPConfig model)
- **clap** -- [`RobertaTokenizer`] (ClapConfig model)
- **clip** -- [`CLIPTokenizer`] (CLIPConfig model)
- **clipseg** -- [`CLIPTokenizer`] (CLIPSegConfig model)
- **clvp** -- [`ClvpTokenizer`] (ClvpConfig model)
- **codegen** -- [`GPT2Tokenizer`] (CodeGenConfig model)
- **cohere** -- [`CohereTokenizer`] (CohereConfig model)
- **cohere2** -- [`CohereTokenizer`] (Cohere2Config model)
- **cohere_asr** -- [`TokenizersBackend`] (CohereAsrConfig model)
- **colqwen2** -- [`Qwen2Tokenizer`] (ColQwen2Config model)
- **convbert** -- [`BertTokenizer`] (ConvBertConfig model)
- **cpmant** -- [`CpmAntTokenizer`] (CpmAntConfig model)
- **ctrl** -- [`CTRLTokenizer`] (CTRLConfig model)
- **data2vec-audio** -- [`Wav2Vec2CTCTokenizer`] (Data2VecAudioConfig model)
- **data2vec-text** -- [`RobertaTokenizer`] (Data2VecTextConfig model)
- **dbrx** -- [`GPT2Tokenizer`] (DbrxConfig model)
- **deberta** -- [`DebertaTokenizer`] (DebertaConfig model)
- **deberta-v2** -- [`DebertaV2Tokenizer`] (DebertaV2Config model)
- **deepseek_v2** -- [`TokenizersBackend`] (DeepseekV2Config model)
- **deepseek_v3** -- [`TokenizersBackend`] (DeepseekV3Config model)
- **deepseek_vl** -- [`TokenizersBackend`] (DeepseekVLConfig model)
- **deepseek_vl_hybrid** -- [`TokenizersBackend`] (DeepseekVLHybridConfig model)
- **dia** -- [`DiaTokenizer`] (DiaConfig model)
- **distilbert** -- [`BertTokenizer`] (DistilBertConfig model)
- **dpr** -- [`DPRQuestionEncoderTokenizer`] (DPRConfig model)
- **electra** -- [`BertTokenizer`] (ElectraConfig model)
- **emu3** -- [`GPT2Tokenizer`] (Emu3Config model)
- **ernie** -- [`BertTokenizer`] (ErnieConfig model)
- **esm** -- [`EsmTokenizer`] (EsmConfig model)
- **falcon_mamba** -- [`GPTNeoXTokenizer`] (FalconMambaConfig model)
- **fastspeech2_conformer** -- [`None`] (FastSpeech2ConformerConfig model)
- **flaubert** -- [`FlaubertTokenizer`] (FlaubertConfig model)
- **flava** -- [`BertTokenizer`] (FlavaConfig model)
- **flex_olmo** -- [`GPT2Tokenizer`] (FlexOlmoConfig model)
- **florence2** -- [`BartTokenizer`] (Florence2Config model)
- **fnet** -- [`FNetTokenizer`] (FNetConfig model)
- **fsmt** -- [`FSMTTokenizer`] (FSMTConfig model)
- **funnel** -- [`FunnelTokenizer`] (FunnelConfig model)
- **fuyu** -- [`TokenizersBackend`] (FuyuConfig model)
- **gemma** -- [`GemmaTokenizer`] (GemmaConfig model)
- **gemma2** -- [`GemmaTokenizer`] (Gemma2Config model)
- **gemma3** -- [`GemmaTokenizer`] (Gemma3Config model)
- **gemma3_text** -- [`GemmaTokenizer`] (Gemma3TextConfig model)
- **gemma3n** -- [`GemmaTokenizer`] (Gemma3nConfig model)
- **gemma3n_text** -- [`GemmaTokenizer`] (Gemma3nTextConfig model)
- **git** -- [`BertTokenizer`] (GitConfig model)
- **glm** -- [`TokenizersBackend`] (GlmConfig model)
- **glm4** -- [`TokenizersBackend`] (Glm4Config model)
- **glm4_moe** -- [`TokenizersBackend`] (Glm4MoeConfig model)
- **glm4_moe_lite** -- [`TokenizersBackend`] (Glm4MoeLiteConfig model)
- **glm4v** -- [`TokenizersBackend`] (Glm4vConfig model)
- **glm4v_moe** -- [`TokenizersBackend`] (Glm4vMoeConfig model)
- **glm_image** -- [`TokenizersBackend`] (GlmImageConfig model)
- **glmasr** -- [`TokenizersBackend`] (GlmAsrConfig model)
- **got_ocr2** -- [`TokenizersBackend`] (GotOcr2Config model)
- **gpt-sw3** -- [`None`] (GPT2Config model)
- **gpt2** -- [`GPT2Tokenizer`] (GPT2Config model)
- **gpt_bigcode** -- [`GPT2Tokenizer`] (GPTBigCodeConfig model)
- **gpt_neo** -- [`GPT2Tokenizer`] (GPTNeoConfig model)
- **gpt_neox** -- [`GPTNeoXTokenizer`] (GPTNeoXConfig model)
- **gpt_neox_japanese** -- [`GPTNeoXJapaneseTokenizer`] (GPTNeoXJapaneseConfig model)
- **gptj** -- [`GPT2Tokenizer`] (GPTJConfig model)
- **granite** -- [`GPT2Tokenizer`] (GraniteConfig model)
- **granitemoe** -- [`GPT2Tokenizer`] (GraniteMoeConfig model)
- **granitemoehybrid** -- [`GPT2Tokenizer`] (GraniteMoeHybridConfig model)
- **granitemoeshared** -- [`GPT2Tokenizer`] (GraniteMoeSharedConfig model)
- **grounding-dino** -- [`BertTokenizer`] (GroundingDinoConfig model)
- **groupvit** -- [`CLIPTokenizer`] (GroupViTConfig model)
- **hubert** -- [`Wav2Vec2CTCTokenizer`] (HubertConfig model)
- **ibert** -- [`RobertaTokenizer`] (IBertConfig model)
- **idefics** -- [`LlamaTokenizer`] (IdeficsConfig model)
- **idefics2** -- [`LlamaTokenizer`] (Idefics2Config model)
- **instructblip** -- [`GPT2Tokenizer`] (InstructBlipConfig model)
- **instructblipvideo** -- [`GPT2Tokenizer`] (InstructBlipVideoConfig model)
- **internvl** -- [`Qwen2Tokenizer`] (InternVLConfig model)
- **jais2** -- [`GPT2Tokenizer`] (Jais2Config model)
- **jamba** -- [`TokenizersBackend`] (JambaConfig model)
- **janus** -- [`TokenizersBackend`] (JanusConfig model)
- **jina_embeddings_v3** -- [`XLMRobertaTokenizer`] (JinaEmbeddingsV3Config model)
- **kosmos-2** -- [`XLMRobertaTokenizer`] (Kosmos2Config model)
- **lasr_ctc** -- [`LasrTokenizer`] (LasrCTCConfig model)
- **lasr_encoder** -- [`LasrTokenizer`] (LasrEncoderConfig model)
- **layoutlm** -- [`BertTokenizer`] (LayoutLMConfig model)
- **layoutlmv2** -- [`LayoutLMv2Tokenizer`] (LayoutLMv2Config model)
- **layoutlmv3** -- [`LayoutLMv3Tokenizer`] (LayoutLMv3Config model)
- **layoutxlm** -- [`LayoutXLMTokenizer`] (LayoutXLMConfig model)
- **led** -- [`LEDTokenizer`] (LEDConfig model)
- **lighton_ocr** -- [`Qwen2TokenizerFast`] (LightOnOcrConfig model)
- **lilt** -- [`RobertaTokenizer`] (LiltConfig model)
- **llava** -- [`TokenizersBackend`] (LlavaConfig model)
- **llava_next** -- [`TokenizersBackend`] (LlavaNextConfig model)
- **longformer** -- [`RobertaTokenizer`] (LongformerConfig model)
- **luke** -- [`LukeTokenizer`] (LukeConfig model)
- **lxmert** -- [`LxmertTokenizer`] (LxmertConfig model)
- **m2m_100** -- [`None`] (M2M100Config model)
- **mamba** -- [`GPTNeoXTokenizer`] (MambaConfig model)
- **mamba2** -- [`GPTNeoXTokenizer`] (Mamba2Config model)
- **marian** -- [`None`] (MarianConfig model)
- **markuplm** -- [`MarkupLMTokenizer`] (MarkupLMConfig model)
- **mbart** -- [`MBartTokenizer`] (MBartConfig model)
- **megatron-bert** -- [`BertTokenizer`] (MegatronBertConfig model)
- **metaclip_2** -- [`XLMRobertaTokenizer`] (MetaClip2Config model)
- **mgp-str** -- [`MgpstrTokenizer`] (MgpstrConfig model)
- **minicpmv4_6** -- [`TokenizersBackend`] (MiniCPMV4_6Config model)
- **minimax_m2** -- [`TokenizersBackend`] (MiniMaxM2Config model)
- **ministral** -- [`TokenizersBackend`] (MinistralConfig model)
- **ministral3** -- [`TokenizersBackend`] (Ministral3Config model)
- **mistral** -- [`TokenizersBackend`] (MistralConfig model)
- **mistral3** -- [`TokenizersBackend`] (Mistral3Config model)
- **mixtral** -- [`TokenizersBackend`] (MixtralConfig model)
- **mm-grounding-dino** -- [`BertTokenizer`] (MMGroundingDinoConfig model)
- **mobilebert** -- [`MobileBertTokenizer`] (MobileBertConfig model)
- **modernbert** -- [`TokenizersBackend`] (ModernBertConfig model)
- **mpnet** -- [`MPNetTokenizer`] (MPNetConfig model)
- **mpt** -- [`GPTNeoXTokenizer`] (MptConfig model)
- **mra** -- [`RobertaTokenizer`] (MraConfig model)
- **mt5** -- [`T5Tokenizer`] (MT5Config model)
- **musicgen** -- [`T5Tokenizer`] (MusicgenConfig model)
- **musicgen_melody** -- [`T5Tokenizer`] (MusicgenMelodyConfig model)
- **mvp** -- [`MvpTokenizer`] (MvpConfig model)
- **nemotron** -- [`TokenizersBackend`] (NemotronConfig model)
- **nllb-moe** -- [`NllbTokenizer`] (NllbMoeConfig model)
- **nomic_bert** -- [`BertTokenizer`] (NomicBertConfig model)
- **nougat** -- [`NougatTokenizer`] (NougatConfig model)
- **nystromformer** -- [`AlbertTokenizer`] (NystromformerConfig model)
- **olmo** -- [`GPTNeoXTokenizer`] (OlmoConfig model)
- **olmo2** -- [`GPTNeoXTokenizer`] (Olmo2Config model)
- **olmo3** -- [`TokenizersBackend`] (Olmo3Config model)
- **olmo_hybrid** -- [`TokenizersBackend`] (OlmoHybridConfig model)
- **olmoe** -- [`GPTNeoXTokenizer`] (OlmoeConfig model)
- **omdet-turbo** -- [`CLIPTokenizer`] (OmDetTurboConfig model)
- **oneformer** -- [`CLIPTokenizer`] (OneFormerConfig model)
- **openai-gpt** -- [`OpenAIGPTTokenizer`] (OpenAIGPTConfig model)
- **opt** -- [`GPT2Tokenizer`] (OPTConfig model)
- **ovis2** -- [`Qwen2Tokenizer`] (Ovis2Config model)
- **owlv2** -- [`CLIPTokenizer`] (Owlv2Config model)
- **owlvit** -- [`CLIPTokenizer`] (OwlViTConfig model)
- **pegasus** -- [`PegasusTokenizer`] (PegasusConfig model)
- **pegasus_x** -- [`PegasusTokenizer`] (PegasusXConfig model)
- **perceiver** -- [`PerceiverTokenizer`] (PerceiverConfig model)
- **phi** -- [`GPT2Tokenizer`] (PhiConfig model)
- **phi3** -- [`TokenizersBackend`] (Phi3Config model)
- **phimoe** -- [`TokenizersBackend`] (PhimoeConfig model)
- **pix2struct** -- [`T5Tokenizer`] (Pix2StructConfig model)
- **pixtral** -- [`TokenizersBackend`] (PixtralVisionConfig model)
- **plbart** -- [`PLBartTokenizer`] (PLBartConfig model)
- **prophetnet** -- [`ProphetNetTokenizer`] (ProphetNetConfig model)
- **qianfan_ocr** -- [`Qwen2Tokenizer`] (QianfanOCRConfig model)
- **qwen2** -- [`Qwen2Tokenizer`] (Qwen2Config model)
- **qwen2_5_omni** -- [`Qwen2Tokenizer`] (Qwen2_5OmniConfig model)
- **qwen2_5_vl** -- [`Qwen2Tokenizer`] (Qwen2_5_VLConfig model)
- **qwen2_audio** -- [`Qwen2Tokenizer`] (Qwen2AudioConfig model)
- **qwen2_moe** -- [`Qwen2Tokenizer`] (Qwen2MoeConfig model)
- **qwen2_vl** -- [`Qwen2Tokenizer`] (Qwen2VLConfig model)
- **qwen3** -- [`Qwen2Tokenizer`] (Qwen3Config model)
- **qwen3_5** -- [`Qwen3_5Tokenizer`] (Qwen3_5Config model)
- **qwen3_5_moe** -- [`Qwen3_5Tokenizer`] (Qwen3_5MoeConfig model)
- **qwen3_moe** -- [`Qwen2Tokenizer`] (Qwen3MoeConfig model)
- **qwen3_next** -- [`Qwen2Tokenizer`] (Qwen3NextConfig model)
- **qwen3_omni_moe** -- [`Qwen2Tokenizer`] (Qwen3OmniMoeConfig model)
- **qwen3_vl** -- [`Qwen2Tokenizer`] (Qwen3VLConfig model)
- **qwen3_vl_moe** -- [`Qwen2Tokenizer`] (Qwen3VLMoeConfig model)
- **rag** -- [`RagTokenizer`] (RagConfig model)
- **recurrent_gemma** -- [`GemmaTokenizer`] (RecurrentGemmaConfig model)
- **reformer** -- [`ReformerTokenizer`] (ReformerConfig model)
- **rembert** -- [`RemBertTokenizer`] (RemBertConfig model)
- **roberta** -- [`RobertaTokenizer`] (RobertaConfig model)
- **roberta-prelayernorm** -- [`RobertaTokenizer`] (RobertaPreLayerNormConfig model)
- **roc_bert** -- [`RoCBertTokenizer`] (RoCBertConfig model)
- **roformer** -- [`RoFormerTokenizer`] (RoFormerConfig model)
- **rwkv** -- [`GPTNeoXTokenizer`] (RwkvConfig model)
- **sam3** -- [`CLIPTokenizer`] (Sam3Config model)
- **sam3_video** -- [`CLIPTokenizer`] (Sam3VideoConfig model)
- **seamless_m4t** -- [`SeamlessM4TTokenizer`] (SeamlessM4TConfig model)
- **seamless_m4t_v2** -- [`SeamlessM4TTokenizer`] (SeamlessM4Tv2Config model)
- **shieldgemma2** -- [`GemmaTokenizer`] (ShieldGemma2Config model)
- **siglip** -- [`None`] (SiglipConfig model)
- **siglip2** -- [`Siglip2Tokenizer`] (Siglip2Config model)
- **speech_to_text** -- [`None`] (Speech2TextConfig model)
- **speecht5** -- [`None`] (SpeechT5Config model)
- **splinter** -- [`SplinterTokenizer`] (SplinterConfig model)
- **squeezebert** -- [`BertTokenizer`] (SqueezeBertConfig model)
- **stablelm** -- [`GPTNeoXTokenizer`] (StableLmConfig model)
- **starcoder2** -- [`GPT2Tokenizer`] (Starcoder2Config model)
- **switch_transformers** -- [`T5Tokenizer`] (SwitchTransformersConfig model)
- **t5** -- [`T5Tokenizer`] (T5Config model)
- **t5gemma** -- [`GemmaTokenizer`] (T5GemmaConfig model)
- **tapas** -- [`TapasTokenizer`] (TapasConfig model)
- **trocr** -- [`XLMRobertaTokenizer`] (TrOCRConfig model)
- **tvp** -- [`BertTokenizer`] (TvpConfig model)
- **udop** -- [`UdopTokenizer`] (UdopConfig model)
- **umt5** -- [`T5Tokenizer`] (UMT5Config model)
- **unispeech** -- [`Wav2Vec2CTCTokenizer`] (UniSpeechConfig model)
- **unispeech-sat** -- [`Wav2Vec2CTCTokenizer`] (UniSpeechSatConfig model)
- **vilt** -- [`BertTokenizer`] (ViltConfig model)
- **vipllava** -- [`TokenizersBackend`] (VipLlavaConfig model)
- **visual_bert** -- [`BertTokenizer`] (VisualBertConfig model)
- **vits** -- [`VitsTokenizer`] (VitsConfig model)
- **voxtral** -- [`TokenizersBackend`] (VoxtralConfig model)
- **voxtral_realtime** -- [`TokenizersBackend`] (VoxtralRealtimeConfig model)
- **wav2vec2** -- [`Wav2Vec2CTCTokenizer`] (Wav2Vec2Config model)
- **wav2vec2-bert** -- [`Wav2Vec2CTCTokenizer`] (Wav2Vec2BertConfig model)
- **wav2vec2-conformer** -- [`Wav2Vec2CTCTokenizer`] (Wav2Vec2ConformerConfig model)
- **whisper** -- [`WhisperTokenizer`] (WhisperConfig model)
- **xclip** -- [`CLIPTokenizer`] (XCLIPConfig model)
- **xglm** -- [`XGLMTokenizer`] (XGLMConfig model)
- **xlm** -- [`XLMTokenizer`] (XLMConfig model)
- **xlm-roberta** -- [`XLMRobertaTokenizer`] (XLMRobertaConfig model)
- **xlm-roberta-xl** -- [`XLMRobertaTokenizer`] (XLMRobertaXLConfig model)
- **xlnet** -- [`XLNetTokenizer`] (XLNetConfig model)
- **xlstm** -- [`GPTNeoXTokenizer`] (xLSTMConfig model)
- **xmod** -- [`XLMRobertaTokenizer`] (XmodConfig model)
- **yoso** -- [`AlbertTokenizer`] (YosoConfig model)

Params -----= pretrained_model_name_or_path (str or os.PathLike): Can be either:

    - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co.
    - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved
      using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
    - a path to a single saved vocabulary file if and only if the tokenizer only requires a
      single vocabulary file (like Bert or XLNet), e.g.: `./my_model_directory/vocab.txt`. (Not
      applicable to all derived classes)

inputs (additional positional arguments, optional): Will be passed along to the Tokenizer __init__() method. config ([PreTrainedConfig], optional) The configuration object used to determine the tokenizer class to instantiate. cache_dir (str or os.PathLike, optional): Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. force_download (bool, optional, defaults to False): Whether or not to force the (re-)download the model weights and configuration files and override the cached versions if they exist. proxies (dict[str, str], optional): A dictionary of proxy servers to use by protocol or endpoint, e.g., {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}. The proxies are used on each request. revision (str, optional, defaults to "main"): The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a git-based system for storing models and other artifacts on huggingface.co, so revision can be any identifier allowed by git. subfolder (str, optional): In case the relevant files are located inside a subfolder of the model repo on huggingface.co (e.g. for facebook/rag-token-base), specify it here. tokenizer_type (str, optional): Tokenizer type to be loaded. backend (str, optional, defaults to "tokenizers"): Backend to use for tokenization. Valid options are: - "tokenizers": Use the HuggingFace tokenizers library backend (default) - "sentencepiece": Use the SentencePiece backend trust_remote_code (bool, optional, defaults to False): Whether or not to allow for custom models defined on the Hub in their own modeling files. This option should only be set to True for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine. kwargs (additional keyword arguments, optional): Will be passed to the Tokenizer __init__() method. Can be used to set special tokens like bos_token, eos_token, unk_token, sep_token, pad_token, cls_token, mask_token, additional_special_tokens. See parameters in the __init__() for more details.

Examples:

>>> from transformers import AutoTokenizer

>>> # Download vocabulary from huggingface.co and cache.
>>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")

>>> # Download vocabulary from huggingface.co (user-uploaded) and cache.
>>> tokenizer = AutoTokenizer.from_pretrained("dbmdz/bert-base-german-cased")

>>> # If vocabulary files are in a directory (e.g. tokenizer was saved using *save_pretrained('./test/saved_model/')*)
>>> # tokenizer = AutoTokenizer.from_pretrained("./test/bert_saved_model/")

>>> # Download vocabulary from huggingface.co and define model-specific arguments
>>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/roberta-base", add_prefix_space=True)

>>> # Explicitly use the tokenizers backend
>>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer", backend="tokenizers")

>>> # Explicitly use the sentencepiece backend
>>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer", backend="sentencepiece")

def register(config_class,
tokenizer_class=None,
slow_tokenizer_class=None,
fast_tokenizer_class=None,
exist_ok=False)

Expand source code

@staticmethod
def register(
    config_class, tokenizer_class=None, slow_tokenizer_class=None, fast_tokenizer_class=None, exist_ok=False
):
    """
    Register a new tokenizer in this mapping.

    Args:
        config_class ([`PreTrainedConfig`]):
            The configuration corresponding to the model to register.
        tokenizer_class: The tokenizer class to register (V5 - preferred parameter).
        slow_tokenizer_class: (Deprecated) The slow tokenizer to register.
        fast_tokenizer_class: (Deprecated) The fast tokenizer to register.
    """
    if tokenizer_class is None:
        # Legacy: prefer fast over slow
        if fast_tokenizer_class is not None:
            tokenizer_class = fast_tokenizer_class
        elif slow_tokenizer_class is not None:
            tokenizer_class = slow_tokenizer_class
        else:
            raise ValueError("You need to pass a `tokenizer_class`")

    for candidate in (slow_tokenizer_class, fast_tokenizer_class, tokenizer_class):
        if candidate is not None:
            REGISTERED_TOKENIZER_CLASSES[candidate.__name__] = candidate

    if slow_tokenizer_class is not None and fast_tokenizer_class is not None:
        REGISTERED_FAST_ALIASES[slow_tokenizer_class.__name__] = fast_tokenizer_class

    TOKENIZER_MAPPING.register(config_class, tokenizer_class, exist_ok=exist_ok)

Register a new tokenizer in this mapping.

Args

-----=

config_class ([PreTrainedConfig]):

The configuration corresponding to the model to register.

tokenizer_class

The tokenizer class to register (V5 - preferred parameter).

slow_tokenizer_class

(Deprecated) The slow tokenizer to register.

fast_tokenizer_class

(Deprecated) The fast tokenizer to register.

class BAMUMoEBlock (config: Any | None = None,
rank: int = 64,
num_experts: int = 2,
top_k: int = 1,
layer_id: int | None = None,
num_layers: int | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
comp_config: Any | None = None,
moe_enabled: bool = False,
svd_rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "bam_umoe", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
class BAMUMoEBlock(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 64,
        num_experts: int = 2,
        top_k: int = 1,
        layer_id: int | None = None,
        num_layers: int | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        moe_enabled: bool = False,  # BAM-FIX-02
        svd_rank: int | None = None,  # BAM-FIX-03
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        **kwargs: Any,
    ) -> None:
        # BAM-FIX-07: config None-safe
        if config is None:

            class _Cfg:
                pass

            config = _Cfg()
            config.hidden_size = kwargs.get("hidden_size", 768)
            config.num_attention_heads = kwargs.get(
                "num_heads", kwargs.get("num_attention_heads", 12)
            )
            config.head_dim = kwargs.get("head_dim", 64)
            config.model_type = kwargs.get("model_type", "")

        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        self.dim = int(
            kwargs.get("hidden_size")
            or getattr(
                config,
                "hidden_size",
                getattr(config, "d_model", getattr(config, "dim", 768)),
            )
        )
        self.hidden_size = self.dim
        self.num_heads = int(
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or getattr(
                config,
                "num_attention_heads",
                getattr(config, "num_heads", getattr(config, "n_heads", 12)),
            )
        )
        self.num_kv_heads = int(
            getattr(
                config,
                "num_key_value_heads",
                getattr(config, "num_kv_heads", self.num_heads),
            )
        )
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = self.num_heads // self.num_kv_heads

        cfg_hd = getattr(config, "head_dim", None)
        self.head_dim = int(
            cfg_hd or kwargs.get("head_dim") or (self.dim // max(1, self.num_heads))
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        self.lora_rank = int(
            getattr(
                config,
                "lora_rank",
                (
                    getattr(comp_config, "lora_rank", lora_rank)
                    if comp_config is not None
                    else lora_rank
                ),
            )
        )
        self.lora_alpha = float(
            getattr(
                config,
                "lora_alpha",
                (
                    getattr(comp_config, "lora_alpha", lora_alpha)
                    if comp_config is not None
                    else lora_alpha
                ),
            )
        )
        self.freeze_base = self.lora_rank > 0

        # BAM-FIX-03: SVD rank for real compression
        default_svd = min(128, self.dim)
        self.svd_rank = int(
            getattr(
                config,
                "svd_rank",
                (
                    getattr(comp_config, "svd_rank", svd_rank)
                    if comp_config is not None
                    else (svd_rank or default_svd)
                ),
            )
        )

        # Decoder detection
        mt = str(getattr(config, "model_type", "") or "").lower()
        is_dec = bool(getattr(config, "is_decoder", False) or getattr(config, "causal", False))
        if any(
            mt.startswith(p)
            for p in (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        if old_attn is not None and hasattr(old_attn, "c_attn"):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        self.num_experts = max(1, int(num_experts))
        self.top_k = min(max(1, int(top_k)), self.num_experts)
        self.moe_enabled = bool(moe_enabled)

        # ---- Q/K/V/O projections ----
        # BAM-FIX-03: use FactorisedLoRALinear with svd_rank
        # If svd_rank >= min(dim, all_head_size), degenerates to dense (no compression)
        Q_rank = min(self.svd_rank, self.dim, self.all_head_size)
        KV_rank = min(self.svd_rank, self.dim, self.all_kv_head_size)
        O_rank = min(self.svd_rank, self.all_head_size, self.dim)

        self.q_lin = FactorisedLoRALinear(
            self.dim,
            self.all_head_size,
            Q_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )
        self.k_lin = FactorisedLoRALinear(
            self.dim,
            self.all_kv_head_size,
            KV_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )
        self.v_lin = FactorisedLoRALinear(
            self.dim,
            self.all_kv_head_size,
            KV_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )
        self.out_lin = FactorisedLoRALinear(
            self.all_head_size,
            self.dim,
            O_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )

        # ---- Optional MoE ----
        if self.moe_enabled:
            self.qkv_moe: UnifiedMoE | None = UnifiedMoE(
                self.dim,
                num_experts=self.num_experts,
                top_k=self.top_k,
                output_dim=self.all_kv_head_size,  # MoE refines V (matches V dim)
                gate_init=kwargs.get("moe_gate_init", -3.0),
            )
        else:
            self.qkv_moe = None

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
        )
        self.dropout = nn.Dropout(dropout_p)

        self.layer_id = layer_id
        self.num_layers = num_layers
        self.comp_config = comp_config
        self.rotary_emb = None

        # BAM-FIX-09: storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

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

    # ---- aliases ----
    @property
    def q_proj(self):
        return self.q_lin

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

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

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

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

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

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

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

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

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

    # ------------------------------------------------------------------
    def copy_weights(self, pretrained_attn: nn.Module) -> None:
        if pretrained_attn is None:
            return

        try:
            for p in pretrained_attn.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        cls_name = pretrained_attn.__class__.__name__
        has_out_lin = any(
            hasattr(pretrained_attn, a)
            for a in ("out_lin", "c_proj", "o_proj", "out_proj", "dense")
        )
        is_nested = False
        if (
            (
                hasattr(pretrained_attn, "self")
                and hasattr(pretrained_attn.self, "query")
                and hasattr(pretrained_attn, "output")
                and hasattr(pretrained_attn.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)

        if hasattr(pretrained_attn, "rotary_emb") and pretrained_attn.rotary_emb is not None:
            self.rotary_emb = pretrained_attn.rotary_emb

        extracted = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused
        if hasattr(pretrained_attn, "c_attn") and pretrained_attn.c_attn is not None:
            c = pretrained_attn.c_attn
            base_c = getattr(c, "base_layer", c)
            cw = getattr(base_c, "weight", getattr(c, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                is_conv = (
                    c.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (cw.ndim == 2 and cw.shape[0] < cw.shape[1])
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )

            if hasattr(pretrained_attn, "c_proj") and pretrained_attn.c_proj is not None:
                o = pretrained_attn.c_proj
                base_o = getattr(o, "base_layer", o)
                ow = getattr(base_o, "weight", getattr(o, "weight", None))
                if ow is not None:
                    ow = ow.data.clone()
                    is_conv_o = (
                        o.__class__.__name__ == "Conv1D"
                        or base_o.__class__.__name__ == "Conv1D"
                        or (ow.ndim == 2 and ow.shape[0] < ow.shape[1])
                    )
                    if is_conv_o:
                        ow = ow.t().contiguous()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # Nested BERT / RoBERTa
        elif hasattr(pretrained_attn, "self") and hasattr(pretrained_attn.self, "query"):
            s = pretrained_attn.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(pretrained_attn, "output") and hasattr(pretrained_attn.output, "dense"):
                d = pretrained_attn.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat modules
        else:
            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 key, names in candidates.items():
                for n in names:
                    if hasattr(pretrained_attn, n):
                        m = getattr(pretrained_attn, n)
                        if m is not None and hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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(pretrained_attn, 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"BAMUMoE.copy_weights: missing {missing_after} from "
                f"{pretrained_attn.__class__.__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

        target_map = {
            "q": self.q_lin,
            "k": self.k_lin,
            "v": self.v_lin,
            "o": self.out_lin,
        }

        with torch.no_grad():
            for proj_key, target in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        # Identity for nested encoder bypass
                        eye = torch.eye(
                            target.out_features,
                            target.in_features,
                            device=target.U.device,
                            dtype=target.U.dtype,
                        )
                        target.set_svd_factors(eye, None)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                # Orient to (out, in)
                in_dim = self.all_head_size if proj_key == "o" else self.dim
                w = _orient(w, in_dim)

                # BAM-FIX-08: shape tolerance
                out_d = min(target.out_features, w.size(0))
                in_d = min(target.in_features, w.size(1))
                if (out_d, in_d) != (target.out_features, target.in_features):
                    padded = torch.zeros(
                        target.out_features,
                        target.in_features,
                        device=w.device,
                        dtype=w.dtype,
                    )
                    padded[:out_d, :in_d] = w[:out_d, :in_d]
                    w = padded

                # Handle GQA pooling for K/V
                if proj_key in ("k", "v"):
                    src_out = w.shape[0]
                    tgt_out = target.out_features
                    if src_out > tgt_out and src_out % tgt_out == 0:
                        rep = src_out // tgt_out
                        w = w.view(tgt_out, rep, -1).mean(dim=1)
                    elif src_out < tgt_out and tgt_out % src_out == 0:
                        rep = tgt_out // src_out
                        w = w.repeat_interleave(rep, dim=0)

                target.set_svd_factors(w, b)
                target.freeze_base_weights()

        # Track compressed bytes
        try:
            for m in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

    # ------------------------------------------------------------------
    def get_aux_loss(self, device: torch.device) -> torch.Tensor:
        """BAM-FIX-04: MoE aux loss for training loop."""
        if self.qkv_moe is None:
            return torch.zeros((), device=device, dtype=torch.float32)
        return self.qkv_moe.get_aux_loss(device)

    # ------------------------------------------------------------------
    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,
    ):
        B, L, D = hidden_states.shape
        dtype_in = hidden_states.dtype

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

        q = (
            self.q_lin(hidden_states)
            .view(B, L, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_lin(kv_input)
            .view(B, S, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_lin(kv_input)
            .view(B, S, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- RoPE ----
        cos = kwargs.get("cos")
        sin = kwargs.get("sin")
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) == 2:
            cos, sin = pos_emb[0], pos_emb[1]

        if cos is not None and sin is not None and not is_cross:
            q, k = apply_rotary_pos_emb(q, k, cos, sin, kwargs.get("position_ids"))

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

        # ---- MoE refinement (only self-attn, only if enabled) ----
        if self.qkv_moe is not None and not is_cross:
            moe_delta_v = self.qkv_moe(hidden_states)
            moe_delta_v = moe_delta_v.view(B, L, self.num_kv_heads, self.head_dim)
            moe_delta_v = moe_delta_v.transpose(1, 2).contiguous()
            if self.kv_repeat > 1:
                moe_delta_v = repeat_kv(moe_delta_v, self.kv_repeat)
            v = (v + 0.1 * moe_delta_v.to(v.dtype)).to(v.dtype)

        # ---- Dense attention (no SDPA shortcut when MoE enabled, to keep mask logic uniform) ----
        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross:
            causal = torch.triu(
                torch.full((L, S), float("-inf"), device=scores.device, dtype=torch.float32),
                diagonal=1 + (S - L),
            )
            scores = scores + causal.unsqueeze(0).unsqueeze(1)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                scores = scores.masked_fill((eff_mask == 0).unsqueeze(1).unsqueeze(2), -1e4)
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(dtype_in)
        probs = torch.nan_to_num(probs, nan=0.0)
        probs = self.dropout(probs)

        # ---- head_mask (BAM-FIX-10) ----
        if head_mask is not None:
            hm = head_mask.to(probs.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            probs = probs * hm

        context = torch.matmul(probs.to(v.dtype), v)
        context = context.transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # ---- Output projection ----
        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_lin(context)

        if out.shape[-1] != D:
            raise RuntimeError(f"BAMUMoE: output dim {out.shape[-1]} != hidden_size {D}")

        return out, (probs if output_attentions else None)

Every 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.

Instance variables

prop c_proj

Expand source code

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

prop dense

Expand source code

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

prop k_proj

Expand source code

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

prop key

Expand source code

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

prop o_proj

Expand source code

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

prop out_proj

Expand source code

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

prop q_proj

Expand source code

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

prop query

Expand source code

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

prop v_proj

Expand source code

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

prop value

Expand source code

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

Methods

def get_aux_loss(self, device: torch.device) ‑> torch.Tensor

Expand source code

def get_aux_loss(self, device: torch.device) -> torch.Tensor:
    """BAM-FIX-04: MoE aux loss for training loop."""
    if self.qkv_moe is None:
        return torch.zeros((), device=device, dtype=torch.float32)
    return self.qkv_moe.get_aux_loss(device)

BAM-FIX-04: MoE aux loss for training loop.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class BaseAttention (config: Any | None = None,
rank: int | None = None,
**kwargs: Any)

Expand source code

class BaseAttention(nn.Module):
    """Every custom attention module must inherit from this."""

    def __init__(self, config: Any | None = None, rank: int | None = None, **kwargs: Any) -> None:
        super().__init__()
        self.config = config
        self.rank = rank
        self.num_heads = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or (
                getattr(
                    config,
                    "num_attention_heads",
                    getattr(config, "num_heads", getattr(config, "n_heads", 12)),
                )
                if config is not None
                else 12
            )
        )
        self.hidden_size = (
            kwargs.get("hidden_size")
            or kwargs.get("dim")
            or kwargs.get("d_model")
            or (
                getattr(
                    config,
                    "hidden_size",
                    getattr(config, "d_model", getattr(config, "dim", 768)),
                )
                if config is not None
                else 768
            )
        )
        self.head_dim = getattr(
            config,
            "d_kv",
            getattr(
                config,
                "head_dim",
                self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
            ),
        )

    @abstractmethod
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        output_attentions=False,
        **kwargs,
    ) -> tuple[Any, Any | None]:
        """Must return (output, attention_weights_or_None)."""
        raise NotImplementedError("Subclasses of BaseAttention must implement forward()")

    def copy_weights(self, old_attn: Any) -> None:
        """Universal fallback: copy pretrained weights from old attention layer."""
        if old_attn is None:
            return
        import torch

        from neuralops.core.adapters import UniversalAdapter

        adapter = UniversalAdapter()

        try:
            w_q, b_q = adapter.get_weight_bias(old_attn, "q")
            w_k, b_k = adapter.get_weight_bias(old_attn, "k")
            w_v, b_v = adapter.get_weight_bias(old_attn, "v")
            w_o, b_o = adapter.get_weight_bias(old_attn, "o")

            with torch.no_grad():
                dim = self.hidden_size
                for q_name in ["W_q", "q_proj", "q_lin", "Wq", "query", "q"]:
                    if hasattr(self, q_name):
                        mod = getattr(self, q_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_q.shape:
                            mod.weight.copy_(w_q)
                            if b_q is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_q)

                for k_name in ["W_k", "k_proj", "k_lin", "Wk", "key", "k"]:
                    if hasattr(self, k_name):
                        mod = getattr(self, k_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_k.shape:
                            mod.weight.copy_(w_k)
                            if b_k is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_k)

                for v_name in ["W_v", "v_proj", "v_lin", "Wv", "value", "v"]:
                    if hasattr(self, v_name):
                        mod = getattr(self, v_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_v.shape:
                            mod.weight.copy_(w_v)
                            if b_v is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_v)

                for o_name in ["out", "out_proj", "o_proj", "out_lin", "proj", "dense"]:
                    if hasattr(self, o_name):
                        mod = getattr(self, o_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_o.shape:
                            mod.weight.copy_(w_o)
                            if b_o is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_o)

                # Handle fused in_proj
                if hasattr(self, "in_proj") and hasattr(self.in_proj, "weight"):
                    if self.in_proj.weight.shape == torch.Size([dim * 2, dim]):
                        self.in_proj.weight[:dim].copy_(w_v[:dim])
                        self.in_proj.weight[dim:].copy_(w_q[:dim])
                    elif self.in_proj.weight.shape == torch.Size([dim * 3, dim]):
                        self.in_proj.weight[:dim].copy_(w_q[:dim])
                        self.in_proj.weight[dim : 2 * dim].copy_(w_k[:dim])
                        self.in_proj.weight[2 * dim :].copy_(w_v[:dim])
        except (AttributeError, RuntimeError, ValueError) as e:
            logger.debug("BaseAttention.copy_weights failed: %s", e)
        except Exception as e:
            logger.warning("Unexpected error in BaseAttention.copy_weights: %s", e)

    def extra_repr(self) -> str:
        """Extra repr."""
        return f"num_heads={self.num_heads}, head_dim={self.head_dim}, rank={self.rank}"

Every custom attention module must inherit from this.

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

Ancestors

• torch.nn.modules.module.Module

Subclasses

• AdaptiveSpanAttention
• AdaptiveSparseAttention
• AdaptiveSVDAttention
• AdaSplashAttention
• AegisHybridAttention
• AegisSpectronFP8Attention
• AegisAttentionV2
• AegisAttention
• AegisAttentionV3
• ASVDAttention
• BAMUMoEBlock
• BlockTransformerHIBA
• ChelaSlayAttention
• ClusterAttention
• CompressedAttentionLayer
• CompressedDistilBertAttention
• CosformerAttention
• DecayLinearAttention
• EATBlock
• ExpertChoiceMoEAttention
• FactorisedMatrixState
• FlashAttention3
• FlashAttention3FP8
• FullAttention
• FusedSVDKernelAttention
• FUSIONATTN
• GalerkinAttention
• GatedDeltaNet
• GatedDeltaNetBlock
• GatedLinearAttention
• GoldFinchAHNBlock
• GroupedQueryAttention
• GraniteHybridBlock
• GriffinSSDBlock
• HALOReGLAAttention
• HammingAttention
• HGRN2GatedDeltaNet
• HGRN2Layer
• HGRN2OuterBlock
• HLAAttention
• HybridAttention
• HymbaTransXBlock
• InfiniAttention
• Jamba2Block
• JambaInfoBlock
• KDAAttention
• LEKAV1Attention
• LEKAv2Attention
• LEKAV3Attention
• LinformerAttention
• LoLAAttention
• LoRAAttention
• LSHAttention
• MambaSVDFlash
• MinimalNOVAAttention
• MinimalSVDAttention
• MixAttentionMoSKA
• MLAAttention
• MoBAAttention
• NASCAttention
• NEXUSAttention
• NativeHybridAttention
• NativeSparseAttention
• NystromAttention
• OptimizedHybridAttention
• ORIONAttention
• PADReAttention
• PerformerAttention
• PHANTOMAttention
• PrimalAttention
• ProgressiveTokenPruning
• QCQAAttention
• RALALinearAttention
• REAPMoEAttention
• ReGLAAttention
• ReplaceAttention
• ResSVDAttention
• RoutingAttention
• RPCAttention
• SALTAttention
• SambaMoEBlock
• SelectiveSSM
• SelectiveSSM
• SOAAttention
• SpanExpandedAttention
• SpectraV2Attention
• SpectraformerAttention
• SpectralMambaAttention
• SpectronAttention
• SqueezingHeadsDistiller
• SVDLinearAttention
• SVDCompressedAttention
• SVDLinearAttention
• SVDLLMV2Attention
• SVDMamba
• SVDProjectionAttention
• YourCustomAttention
• TLinFormerAttention
• TokenCompressedAttention
• TPAttention
• TritonCosformerAttention
• TritonSVDLinearAttention
• UnifiedRoPE
• SelectiveMamba
• ViGAttention
• VSpecAttention
• WeightedGroupedQueryAttention
• XAttention
• ZeroSAttention
• NLISoftmaxAttention
• ReplaceAttention

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Universal fallback: copy pretrained weights from old attention layer."""
    if old_attn is None:
        return
    import torch

    from neuralops.core.adapters import UniversalAdapter

    adapter = UniversalAdapter()

    try:
        w_q, b_q = adapter.get_weight_bias(old_attn, "q")
        w_k, b_k = adapter.get_weight_bias(old_attn, "k")
        w_v, b_v = adapter.get_weight_bias(old_attn, "v")
        w_o, b_o = adapter.get_weight_bias(old_attn, "o")

        with torch.no_grad():
            dim = self.hidden_size
            for q_name in ["W_q", "q_proj", "q_lin", "Wq", "query", "q"]:
                if hasattr(self, q_name):
                    mod = getattr(self, q_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_q.shape:
                        mod.weight.copy_(w_q)
                        if b_q is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_q)

            for k_name in ["W_k", "k_proj", "k_lin", "Wk", "key", "k"]:
                if hasattr(self, k_name):
                    mod = getattr(self, k_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_k.shape:
                        mod.weight.copy_(w_k)
                        if b_k is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_k)

            for v_name in ["W_v", "v_proj", "v_lin", "Wv", "value", "v"]:
                if hasattr(self, v_name):
                    mod = getattr(self, v_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_v.shape:
                        mod.weight.copy_(w_v)
                        if b_v is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_v)

            for o_name in ["out", "out_proj", "o_proj", "out_lin", "proj", "dense"]:
                if hasattr(self, o_name):
                    mod = getattr(self, o_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_o.shape:
                        mod.weight.copy_(w_o)
                        if b_o is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_o)

            # Handle fused in_proj
            if hasattr(self, "in_proj") and hasattr(self.in_proj, "weight"):
                if self.in_proj.weight.shape == torch.Size([dim * 2, dim]):
                    self.in_proj.weight[:dim].copy_(w_v[:dim])
                    self.in_proj.weight[dim:].copy_(w_q[:dim])
                elif self.in_proj.weight.shape == torch.Size([dim * 3, dim]):
                    self.in_proj.weight[:dim].copy_(w_q[:dim])
                    self.in_proj.weight[dim : 2 * dim].copy_(w_k[:dim])
                    self.in_proj.weight[2 * dim :].copy_(w_v[:dim])
    except (AttributeError, RuntimeError, ValueError) as e:
        logger.debug("BaseAttention.copy_weights failed: %s", e)
    except Exception as e:
        logger.warning("Unexpected error in BaseAttention.copy_weights: %s", e)

Universal fallback: copy pretrained weights from old attention layer.

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    """Extra repr."""
    return f"num_heads={self.num_heads}, head_dim={self.head_dim}, rank={self.rank}"

Extra repr.

def forward(self,
hidden_states,
attention_mask=None,
head_mask=None,
output_attentions=False,
**kwargs) ‑> tuple[typing.Any, Any | None]

Expand source code

@abstractmethod
def forward(
    self,
    hidden_states,
    attention_mask=None,
    head_mask=None,
    output_attentions=False,
    **kwargs,
) -> tuple[Any, Any | None]:
    """Must return (output, attention_weights_or_None)."""
    raise NotImplementedError("Subclasses of BaseAttention must implement forward()")

Must return (output, attention_weights_or_None).

class BlockTransformerHIBA (config: Any | None = None,
rank: int = 64,
block_size: int = 64,
top_k_inter: int = 4,
rectification_interval: int = 4,
layer_id: int | None = None,
layer_idx: int | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("block_transformer_hiba")
class BlockTransformerHIBA(BaseAttention):
    r"""
    Block Transformer with Hierarchical Intra-Block and Inter-Block Attention (HIBA).

    Hardened for production use across DistilBERT, RoBERTa/BERT, GPT-2, and Qwen/LLaMA/Mistral.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 64,
        block_size: int = 64,
        top_k_inter: int = 4,
        rectification_interval: int = 4,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        old_attn: nn.Module | None = None,
        pretrained_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        source_attn = (
            old_attn
            if old_attn is not None
            else (
                pretrained_attn
                if pretrained_attn is not None
                else kwargs.get("old_attn", kwargs.get("pretrained_attn"))
            )
        )

        self.dim: int = kwargs.get("hidden_size") or getattr(
            config,
            "hidden_size",
            getattr(config, "d_model", getattr(config, "dim", 768)),
        )
        self.hidden_size: int = self.dim
        self.num_heads: int = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or getattr(
                config,
                "num_attention_heads",
                getattr(config, "num_heads", getattr(config, "n_heads", 12)),
            )
        )
        self.num_kv_heads: int = getattr(
            config,
            "num_key_value_heads",
            getattr(config, "num_kv_heads", self.num_heads),
        )
        self.head_dim: int = getattr(
            config,
            "head_dim",
            getattr(config, "kv_channels", self.dim // max(1, self.num_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
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa: bool = True

        self.block_size: int = block_size
        self.top_k_inter: int = top_k_inter
        self.rectification_interval: int = max(1, rectification_interval)

        # Static layer index (not mutated inside forward)
        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx: int = int(resolved_layer_id)

        # LoRA adaptation configuration
        self.lora_rank: int = getattr(
            config,
            "lora_rank",
            getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha: float = getattr(
            config,
            "lora_alpha",
            getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base: bool = self.lora_rank > 0

        # Decoder & nested encoder detection
        model_type = getattr(config, "model_type", "").lower()
        is_dec = getattr(
            config,
            "is_decoder",
            getattr(config, "is_causal", kwargs.get("is_causal", False)),
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(_mt.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)
        self.is_nested_encoder: bool = False

        # Multi-Head Attention Projections equipped with trainable residual LoRA adapters
        self.W_q = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_k = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_v = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        self.rotary_emb = None
        self.comp_config: Any | None = comp_config

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

    # --------------------------------------------------------------------------
    # Compatibility properties for parameter inspection & framework integration
    # --------------------------------------------------------------------------
    @property
    def q_proj(self) -> LoRALinear:
        return self.W_q

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

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

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

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

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

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

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

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

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

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

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

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

    def copy_weights(self, pretrained_attn: nn.Module) -> None:
        r"""
        Extracts and copies pretrained weights across DistilBERT, RoBERTa/BERT, GPT-2, and Qwen/LLaMA.
        Strictly validates that copied weights are non-empty and verifies reconstruction norms.
        """
        if pretrained_attn is None:
            return

        # 1. Detect nested encoder architectures (RoBERTa / BERT)
        cls_name = pretrained_attn.__class__.__name__
        has_out_lin = hasattr(pretrained_attn, "out_lin")
        is_nested = False
        if (
            (
                hasattr(pretrained_attn, "self")
                and hasattr(pretrained_attn.self, "query")
                and hasattr(pretrained_attn, "output")
                and hasattr(pretrained_attn.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)

        # Inherit rotary embeddings if present
        if hasattr(pretrained_attn, "rotary_emb"):
            self.rotary_emb = pretrained_attn.rotary_emb

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 2. GPT-2: Fused c_attn + c_proj
        if hasattr(pretrained_attn, "c_attn") and pretrained_attn.c_attn is not None:
            c_attn_mod = pretrained_attn.c_attn
            if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    print(f"copy_weights failed: {e}")
                    logger.warning(f"copy_weights failed: {e}")
            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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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,
                )

            # GPT-2 Output: c_proj
            if hasattr(pretrained_attn, "c_proj") and pretrained_attn.c_proj is not None:
                c_proj_mod = pretrained_attn.c_proj
                if hasattr(c_proj_mod, "merge") and not getattr(c_proj_mod, "merged", False):
                    try:
                        c_proj_mod.merge()
                    except Exception as e:
                        print(f"copy_weights failed: {e}")
                        logger.warning(f"copy_weights failed: {e}")
                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()
                    if (
                        (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])
                    ):
                        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,
                    )

        # 3. RoBERTa / BERT nested self-attention (pretrained_attn.self)
        elif hasattr(pretrained_attn, "self") and hasattr(pretrained_attn.self, "query"):
            s = pretrained_attn.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,
                    )
            # Output dense from outer container
            if hasattr(pretrained_attn, "output") and hasattr(pretrained_attn.output, "dense"):
                out_mod = pretrained_attn.output.dense
                extracted["o"] = (
                    out_mod.weight.data.clone(),
                    out_mod.bias.data.clone() if out_mod.bias is not None else None,
                )

        # 4. Standard linear projections (DistilBERT, Qwen, LLaMA)
        else:
            proj_candidates = {
                "q": ["q_lin", "q_proj", "query", "W_q"],
                "k": ["k_lin", "k_proj", "key", "W_k"],
                "v": ["v_lin", "v_proj", "value", "W_v"],
                "o": ["out_lin", "out_proj", "dense", "o_proj"],
            }
            for proj_key, names in proj_candidates.items():
                for name in names:
                    if hasattr(pretrained_attn, name):
                        mod = getattr(pretrained_attn, 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

        # 5. Fallback via UniversalAdapter
        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(pretrained_attn, k)
                    if w is not None:
                        extracted[k] = (w.clone(), b.clone() if b is not None else None)
            except Exception:
                pass

        # Verify all mandatory projections were extracted
        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:
            logger.warning(
                f"BlockTransformerHIBA copy_weights failed: missing mandatory projections {missing_after} "
                f"from module {pretrained_attn.__class__.__name__} ({pretrained_attn}). Proceeding with initialized weights."
            )

        # 6. Copy weights into target LoRALinear layers with strict dimension safety
        target_map = {
            "q": self.W_q,
            "k": self.W_k,
            "v": self.W_v,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        # Inner nested self-attention bypasses out_proj; initialize to identity
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                # Handle transposed orientation if necessary (only if shape does not match already)
                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                # Weight verification: ensure weights are non-trivial
                w_norm = float(torch.norm(target_mod.weight.data).item())
                if w_norm < 1e-4:
                    logger.warning(
                        f"Weight copy verification failed for projection '{proj_key}': norm too small ({w_norm:.2e})"
                    )

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    def _convert_mask_to_bool(
        self,
        attention_mask: torch.Tensor | None,
        B: int,
        L: int,
        device: torch.device,
    ) -> torch.Tensor:
        if attention_mask is None:
            return torch.ones(B, L, device=device, dtype=torch.bool)
        if attention_mask.dim() == 2:
            return (
                attention_mask.bool() if attention_mask.dtype == torch.bool else attention_mask > 0
            )
        elif attention_mask.dim() == 3:
            if attention_mask.shape[1] == 1:
                m = attention_mask[:, 0, :]
                return (
                    m > -1000.0
                    if (m.dtype != torch.bool and m.min() < -100.0)
                    else (m if m.dtype == torch.bool else m > 0)
                )
            else:
                m = (
                    attention_mask > -1000.0
                    if (attention_mask.dtype != torch.bool and attention_mask.min() < -100.0)
                    else (
                        attention_mask if attention_mask.dtype == torch.bool else attention_mask > 0
                    )
                )
                return m.any(dim=1)
        elif attention_mask.dim() == 4:
            m = (
                attention_mask > -1000.0
                if (attention_mask.dtype != torch.bool and attention_mask.min() < -100.0)
                else (attention_mask if attention_mask.dtype == torch.bool else attention_mask > 0)
            )
            if m.shape[1] == 1:
                m = m.squeeze(1)
            else:
                m = m.any(dim=1)
            return m.any(dim=-2)
        return torch.ones(B, L, device=device, dtype=torch.bool)

    def _intra_block_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        key_padding_mask: torch.Tensor,
        causal: bool = False,
        head_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, L_b, D_h = q.shape
        scale = 1.0 / math.sqrt(D_h)
        attn = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        if causal:
            causal_mask = torch.triu(
                torch.full((L_b, L_b), -10000.0, device=q.device, dtype=torch.float32),
                diagonal=1,
            )
            attn = attn + causal_mask.unsqueeze(0).unsqueeze(1)

        pad_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
        attn = attn.masked_fill(pad_mask, -10000.0)

        probs = F.softmax(attn, dim=-1)
        probs = torch.nan_to_num(probs, nan=0.0, posinf=1.0, neginf=0.0)
        probs = self.dropout(probs)

        if head_mask is not None:
            probs = probs * head_mask

        return torch.matmul(probs.to(v.dtype), v)

    def _inter_block_sparse(
        self,
        q_b: torch.Tensor,
        k_all: torch.Tensor,
        v_all: torch.Tensor,
        block_idx: int,
        num_blocks: int,
        valid_mask: torch.Tensor,
        head_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, L_b, D_h = q_b.shape
        total_len = k_all.size(2)
        start_b = block_idx * self.block_size
        end_b = min((block_idx + 1) * self.block_size, total_len)

        k_select: list[torch.Tensor] = []
        v_select: list[torch.Tensor] = []
        valid_select: list[torch.Tensor] = []

        # Previous tokens: includes block 0 (CLS at index 0)
        if block_idx > 0:
            k_select.append(k_all[:, :, :start_b, :])
            v_select.append(v_all[:, :, :start_b, :])
            valid_select.append(valid_mask[:, :start_b])

        # Subsequent tokens (only for bidirectional non-causal encoders)
        if (block_idx + 1) * self.block_size < total_len and not self.is_decoder:
            k_select.append(k_all[:, :, end_b:, :])
            v_select.append(v_all[:, :, end_b:, :])
            valid_select.append(valid_mask[:, end_b:])

        if not k_select:
            return torch.zeros_like(q_b)

        k_neighbors = torch.cat(k_select, dim=2)
        v_neighbors = torch.cat(v_select, dim=2)
        valid_neighbors = torch.cat(valid_select, dim=1)

        scale = 1.0 / math.sqrt(D_h)
        scores = (
            torch.matmul(q_b.float(), k_neighbors.float().transpose(-2, -1).contiguous()) * scale
        )

        pad_mask = ~valid_neighbors.unsqueeze(1).unsqueeze(2)
        scores = scores.masked_fill(pad_mask, -10000.0)

        # Guarantee CLS block inclusion for block_idx > 0
        scores_for_topk = scores.clone()
        if block_idx > 0:
            scores_for_topk[:, :, :, 0] = scores_for_topk[:, :, :, 0] + 1e4

        top_k = min(self.top_k_inter, k_neighbors.size(2))
        _, topk_idx = torch.topk(scores_for_topk, top_k, dim=-1)
        sparse_mask = torch.zeros_like(scores).scatter_(-1, topk_idx, 1.0)
        scores = scores.masked_fill(sparse_mask == 0, -10000.0)

        probs = F.softmax(scores, dim=-1)
        probs = torch.nan_to_num(probs, nan=0.0, posinf=1.0, neginf=0.0)
        probs = self.dropout(probs)

        if head_mask is not None:
            probs = probs * head_mask

        return torch.matmul(probs.to(v_neighbors.dtype), v_neighbors)

    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]:
        B, L, D = hidden_states.shape
        H = self.num_heads
        is_cross_attention = encoder_hidden_states is not None
        kv_states = encoder_hidden_states if is_cross_attention else hidden_states
        S = kv_states.shape[1]

        # 1. Base Projections
        q = self.W_q(hidden_states).view(B, L, H, self.head_dim).transpose(1, 2).contiguous()
        k = (
            self.W_k(kv_states)
            .view(B, S, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.W_v(kv_states)
            .view(B, S, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. RoPE Embeddings applied before GQA expansion
        pos_emb = kwargs.get("position_embeddings")
        cos = kwargs.get(
            "cos",
            (pos_emb[0] if isinstance(pos_emb, (tuple, list)) and len(pos_emb) == 2 else None),
        )
        sin = kwargs.get(
            "sin",
            (pos_emb[1] if isinstance(pos_emb, (tuple, list)) and len(pos_emb) == 2 else None),
        )

        if cos is not None and sin is not None:
            q, k = apply_rotary_pos_emb(q, k, cos, sin, kwargs.get("position_ids"))
        elif hasattr(self, "rotary_emb") and self.rotary_emb is not None:
            try:
                from neuralops.utils.compat import maybe_apply_rotary

                q, k = maybe_apply_rotary(
                    q,
                    k,
                    getattr(self, "adapter", None),
                    self.rotary_emb,
                    kwargs.get("position_ids"),
                )
            except Exception:
                pass

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

        # 4. Safe head mask reshaping
        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.shape[0], head_mask.shape[1], 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)

        # Valid token mask
        active_mask = encoder_attention_mask if is_cross_attention else attention_mask
        valid_mask = self._convert_mask_to_bool(active_mask, B, S, hidden_states.device)

        num_blocks = (L + self.block_size - 1) // self.block_size
        # Static rectification decision (does not increment layer_idx across forward passes)
        use_rectify = (
            (self.layer_idx % self.rectification_interval == 0)
            or (num_blocks <= 2)
            or is_cross_attention
        )

        if use_rectify:
            # Dense attention pathway (full rectification)
            scale = 1.0 / math.sqrt(self.head_dim)
            # Memory-efficient: use SDPA to avoid OOM on large seqs

            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 = getattr(self, "is_decoder", False) and not is_cross_attention and L > 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, L, self.all_head_size)
                )
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                return (
                    self.out_proj(_sdpa_out)
                    if not getattr(self, "is_nested_encoder", False)
                    else _sdpa_out
                ), None
            except Exception:
                pass

            scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

            # If a 4D attention mask is provided, it already incorporates causal + padding constraints
            if active_mask is not None and active_mask.dim() == 4:
                if active_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~active_mask, -10000.0)
                else:
                    scores = scores + active_mask.to(scores.dtype)
            else:
                # Causal mask for decoders
                if self.is_decoder and not is_cross_attention:
                    causal_mask = torch.triu(
                        torch.full((L, S), -10000.0, device=scores.device, dtype=torch.float32),
                        diagonal=1 + (S - L),
                    )
                    scores = scores + causal_mask.unsqueeze(0).unsqueeze(1)

                # Padding mask
                pad_mask = ~valid_mask.unsqueeze(1).unsqueeze(2)
                scores = scores.masked_fill(pad_mask, -10000.0)

            attn_probs = F.softmax(scores, dim=-1)
            attn_probs = torch.nan_to_num(attn_probs, nan=0.0, posinf=1.0, neginf=0.0)
            attn_probs = self.dropout(attn_probs)

            if head_mask is not None:
                attn_probs = attn_probs * head_mask

            out = torch.matmul(attn_probs.to(v.dtype), v)
            return_attn = attn_probs if output_attentions else None
        else:
            # Hierarchical Intra-Block and Inter-Block Attention
            out_list = []
            for b_idx in range(num_blocks):
                start = b_idx * self.block_size
                end = min(start + self.block_size, L)
                q_b = q[:, :, start:end, :]
                k_b = k[:, :, start:end, :]
                v_b = v[:, :, start:end, :]
                block_valid = valid_mask[:, start:end]
                block_pad_mask = ~block_valid

                intra_out = self._intra_block_attention(
                    q_b,
                    k_b,
                    v_b,
                    block_pad_mask,
                    causal=(self.is_decoder and not is_cross_attention),
                    head_mask=head_mask,
                )
                inter_out = self._inter_block_sparse(
                    q_b, k, v, b_idx, num_blocks, valid_mask, head_mask=head_mask
                )
                out_b = intra_out + inter_out
                out_list.append(out_b)

            out = torch.cat(out_list, dim=2)
            return_attn = None

        # Merge heads: (B, H, L, head_dim) -> (B, L, all_head_size)
        out = out.transpose(1, 2).contiguous().view(B, L, self.all_head_size)

        # 5. Output projection with nested encoder bypass
        if self.is_nested_encoder:
            # Outer RobertaAttention.output.dense performs the dense projection; bypass out_proj
            final_out = out
        else:
            final_out = self.out_proj(out)

        return final_out, return_attn

Block Transformer with Hierarchical Intra-Block and Inter-Block Attention (HIBA).

Hardened for production use across DistilBERT, RoBERTa/BERT, GPT-2, and Qwen/LLaMA/Mistral.

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.

Instance variables

prop c_proj : LoRALinear

Expand source code

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

prop dense : LoRALinear

Expand source code

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

prop k_lin : LoRALinear

Expand source code

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

prop k_proj : LoRALinear

Expand source code

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

prop key : LoRALinear

Expand source code

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

prop o_proj : LoRALinear

Expand source code

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

prop out_lin : LoRALinear

Expand source code

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

prop q_lin : LoRALinear

Expand source code

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

prop q_proj : LoRALinear

Expand source code

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

prop query : LoRALinear

Expand source code

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

prop v_lin : LoRALinear

Expand source code

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

prop v_proj : LoRALinear

Expand source code

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

prop value : LoRALinear

Expand source code

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

Methods

def copy_weights(self, pretrained_attn: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, pretrained_attn: nn.Module) -> None:
    r"""
    Extracts and copies pretrained weights across DistilBERT, RoBERTa/BERT, GPT-2, and Qwen/LLaMA.
    Strictly validates that copied weights are non-empty and verifies reconstruction norms.
    """
    if pretrained_attn is None:
        return

    # 1. Detect nested encoder architectures (RoBERTa / BERT)
    cls_name = pretrained_attn.__class__.__name__
    has_out_lin = hasattr(pretrained_attn, "out_lin")
    is_nested = False
    if (
        (
            hasattr(pretrained_attn, "self")
            and hasattr(pretrained_attn.self, "query")
            and hasattr(pretrained_attn, "output")
            and hasattr(pretrained_attn.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)

    # Inherit rotary embeddings if present
    if hasattr(pretrained_attn, "rotary_emb"):
        self.rotary_emb = pretrained_attn.rotary_emb

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

    # 2. GPT-2: Fused c_attn + c_proj
    if hasattr(pretrained_attn, "c_attn") and pretrained_attn.c_attn is not None:
        c_attn_mod = pretrained_attn.c_attn
        if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
            try:
                c_attn_mod.merge()
            except Exception as e:
                print(f"copy_weights failed: {e}")
                logger.warning(f"copy_weights failed: {e}")
        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()
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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,
            )

        # GPT-2 Output: c_proj
        if hasattr(pretrained_attn, "c_proj") and pretrained_attn.c_proj is not None:
            c_proj_mod = pretrained_attn.c_proj
            if hasattr(c_proj_mod, "merge") and not getattr(c_proj_mod, "merged", False):
                try:
                    c_proj_mod.merge()
                except Exception as e:
                    print(f"copy_weights failed: {e}")
                    logger.warning(f"copy_weights failed: {e}")
            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()
                if (
                    (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])
                ):
                    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,
                )

    # 3. RoBERTa / BERT nested self-attention (pretrained_attn.self)
    elif hasattr(pretrained_attn, "self") and hasattr(pretrained_attn.self, "query"):
        s = pretrained_attn.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,
                )
        # Output dense from outer container
        if hasattr(pretrained_attn, "output") and hasattr(pretrained_attn.output, "dense"):
            out_mod = pretrained_attn.output.dense
            extracted["o"] = (
                out_mod.weight.data.clone(),
                out_mod.bias.data.clone() if out_mod.bias is not None else None,
            )

    # 4. Standard linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q": ["q_lin", "q_proj", "query", "W_q"],
            "k": ["k_lin", "k_proj", "key", "W_k"],
            "v": ["v_lin", "v_proj", "value", "W_v"],
            "o": ["out_lin", "out_proj", "dense", "o_proj"],
        }
        for proj_key, names in proj_candidates.items():
            for name in names:
                if hasattr(pretrained_attn, name):
                    mod = getattr(pretrained_attn, 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

    # 5. Fallback via UniversalAdapter
    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(pretrained_attn, k)
                if w is not None:
                    extracted[k] = (w.clone(), b.clone() if b is not None else None)
        except Exception:
            pass

    # Verify all mandatory projections were extracted
    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:
        logger.warning(
            f"BlockTransformerHIBA copy_weights failed: missing mandatory projections {missing_after} "
            f"from module {pretrained_attn.__class__.__name__} ({pretrained_attn}). Proceeding with initialized weights."
        )

    # 6. Copy weights into target LoRALinear layers with strict dimension safety
    target_map = {
        "q": self.W_q,
        "k": self.W_k,
        "v": self.W_v,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    # Inner nested self-attention bypasses out_proj; initialize to identity
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            # Handle transposed orientation if necessary (only if shape does not match already)
            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device, dtype=target_mod.weight.dtype
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

            # Weight verification: ensure weights are non-trivial
            w_norm = float(torch.norm(target_mod.weight.data).item())
            if w_norm < 1e-4:
                logger.warning(
                    f"Weight copy verification failed for projection '{proj_key}': norm too small ({w_norm:.2e})"
                )

            if self.freeze_base:
                target_mod.freeze_base_weights()

Extracts and copies pretrained weights across DistilBERT, RoBERTa/BERT, GPT-2, and Qwen/LLaMA. Strictly validates that copied weights are non-empty and verifies reconstruction norms.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class ChelaSlayAttention (config: Any | None = None,
rank: int = 64,
feature_dim: int = 128,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
gate_init_bias: float = 5.0,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "chela_slay", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "chela_slay_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class ChelaSlayAttention(BaseAttention):
    r"""
    CHELA-SLAY: Spherical Linearized Attention with Yat-Kernel.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 64,
        feature_dim: int = 128,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        gate_init_bias: float = 5.0,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config=config, rank=rank, **kwargs)
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # CS-FIX-17: resolve config from source_attn if config is missing
        if config is None and source_attn is not None:

            class _Dummy:
                pass

            config = _Dummy()
            out_layer = (
                _safe_getattr(source_attn, "out_lin", None)
                or _safe_getattr(source_attn, "dense", None)
                or _safe_getattr(source_attn, "c_proj", None)
                or _safe_getattr(source_attn, "o_proj", None)
            )
            q_layer = (
                _safe_getattr(source_attn, "q_lin", None)
                or _safe_getattr(source_attn, "query", None)
                or _safe_getattr(source_attn, "c_attn", None)
                or _safe_getattr(source_attn, "q_proj", None)
            )
            if out_layer is not None and hasattr(out_layer, "out_features"):
                config.hidden_size = out_layer.out_features
            elif q_layer is not None and hasattr(q_layer, "in_features"):
                config.hidden_size = q_layer.in_features
            else:
                config.hidden_size = 768
            config.num_attention_heads = (
                _safe_getattr(source_attn, "num_heads", None)
                or _safe_getattr(source_attn, "num_attention_heads", None)
                or _safe_getattr(source_attn, "n_heads", None)
                or 12
            )
            config.num_key_value_heads = config.num_attention_heads
            config.is_decoder = False
            config.model_type = ""

        if config is None:

            class _DefaultCfg:
                pass

            config = _DefaultCfg()
            config.hidden_size = 768
            config.num_attention_heads = 12
            config.num_key_value_heads = 12
            config.is_decoder = False
            config.model_type = ""

        # CS-FIX-16: safe extraction
        self.dim = _safe_int(
            kwargs.get("hidden_size")
            or kwargs.get("d_model")
            or _resolve_attr(config, "hidden_size", "d_model", "dim"),
            768,
        )
        if self.dim <= 0:
            self.dim = 768

        n_heads_raw = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or _resolve_attr(config, "num_attention_heads", "num_heads", "n_heads", "n_head")
        )
        self.num_heads = _safe_int(n_heads_raw, 12)
        if self.num_heads <= 0:
            self.num_heads = 12

        kv_heads_raw = (
            kwargs.get("num_kv_heads")
            or kwargs.get("num_key_value_heads")
            or _resolve_attr(config, "num_key_value_heads", "num_kv_heads")
        )
        self.num_kv_heads = _safe_int(kv_heads_raw, self.num_heads)
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))

        head_dim_raw = kwargs.get("head_dim") or _safe_getattr(config, "head_dim", None)
        if head_dim_raw is not None:
            self.head_dim = _safe_int(head_dim_raw, self.dim // max(1, self.num_heads))
        else:
            self.head_dim = self.dim // max(1, self.num_heads)
        if self.head_dim <= 0:
            self.head_dim = 64

        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.feature_dim = max(2, _safe_int(feature_dim, 128))
        self.output_feature_dim = 2 * self.feature_dim

        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx = _safe_int(resolved_layer_id, 0)

        # LoRA config (CS-FIX-22: comp_config may be None)
        self.lora_rank = _safe_int(
            _safe_getattr(config, "lora_rank", None)
            or _safe_getattr(comp_config, "lora_rank", None)
            or kwargs.get("lora_rank", 16),
            16,
        )
        self.lora_alpha = _safe_float(
            _safe_getattr(config, "lora_alpha", None)
            or _safe_getattr(comp_config, "lora_alpha", None)
            or kwargs.get("lora_alpha", 32.0),
            32.0,
        )
        self.freeze_base = self.lora_rank > 0

        # Decoder detection
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            model_type = str(_safe_getattr(config, "model_type", "") or "").lower()
            is_dec = bool(
                _safe_getattr(config, "is_decoder", False)
                or _safe_getattr(config, "is_causal", False)
                or kwargs.get("is_causal", False)
            )
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            if any(model_type.startswith(p) for p in _dec_pfx) or any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (_safe_getattr(config, "architectures", []) or [])
            ):
                is_dec = True
            if source_attn is not None:
                if _safe_getattr(source_attn, "is_decoder", False) or hasattr(
                    source_attn, "c_attn"
                ):
                    is_dec = True
            self.is_decoder = is_dec
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # Projections
        self.W_q = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_k = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_v = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Yat projection
        self.yat_proj = nn.Parameter(
            torch.randn(self.head_dim, self.feature_dim, dtype=torch.float32) * 0.02
        )

        # Adaptive gate (CS-FIX-24)
        self.lora_gate = nn.Linear(self.dim, self.num_heads)
        nn.init.zeros_(self.lora_gate.weight)
        with torch.no_grad():
            self.lora_gate.bias.data.fill_(float(gate_init_bias))

        dropout_p = _safe_float(
            _resolve_attr(config, "dropout", "attention_probs_dropout_prob") or 0.1,
            0.1,
        )
        self.dropout = nn.Dropout(dropout_p)

        self.rotary_emb = None
        self.comp_config = comp_config

        # Storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

        if source_attn is not None:
            # CS-FIX-18: never raise on failure; only log
            try:
                self._copy_pretrained(source_attn)
            except Exception as e:
                logger.warning(
                    "ChelaSlayAttention._copy_pretrained failed on %s: %s. "
                    "Using initialized weights.",
                    source_attn.__class__.__name__,
                    e,
                )

    # ---- aliases (CS-FIX-23) ----
    @property
    def gate(self):
        return self.lora_gate

    @property
    def random_proj(self):
        return self.yat_proj

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

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

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

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

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

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

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

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

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

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

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

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

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

    # ------------------------------------------------------------------
    # Weight transfer (CS-FIX-20, CS-FIX-21)
    # ------------------------------------------------------------------
    def _copy_pretrained(self, pretrained_attn: nn.Module) -> None:
        if pretrained_attn is None:
            return

        try:
            for p in pretrained_attn.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        cls_name = pretrained_attn.__class__.__name__
        has_out_lin = (
            hasattr(pretrained_attn, "out_lin")
            or hasattr(pretrained_attn, "c_proj")
            or hasattr(pretrained_attn, "o_proj")
        )
        is_nested = False
        if (
            (
                hasattr(pretrained_attn, "self")
                and hasattr(pretrained_attn.self, "query")
                and hasattr(pretrained_attn, "output")
                and hasattr(pretrained_attn.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)

        if hasattr(pretrained_attn, "rotary_emb"):
            self.rotary_emb = pretrained_attn.rotary_emb

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        def _unwrap(mod):
            return getattr(mod, "base_layer", mod)

        # GPT-2 fused
        if hasattr(pretrained_attn, "c_attn") and pretrained_attn.c_attn is not None:
            c_attn_mod = _unwrap(pretrained_attn.c_attn)
            c_w = getattr(c_attn_mod, "weight", None)
            if c_w is not None:
                c_w = c_w.data.clone()
                is_conv = pretrained_attn.c_attn.__class__.__name__ == "Conv1D" or (
                    c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1]
                )
                if is_conv:
                    c_w = c_w.t().contiguous()
                split = c_w.size(0) // 3
                c_b = getattr(c_attn_mod, "bias", None)
                c_b_data = c_b.data.clone() if c_b is not None else None
                extracted["q"] = (
                    c_w[:split],
                    c_b_data[:split] if c_b_data is not None else None,
                )
                extracted["k"] = (
                    c_w[split : 2 * split],
                    c_b_data[split : 2 * split] if c_b_data is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :],
                    c_b_data[2 * split :] if c_b_data is not None else None,
                )

            if hasattr(pretrained_attn, "c_proj") and pretrained_attn.c_proj is not None:
                c_proj_mod = _unwrap(pretrained_attn.c_proj)
                o_w = getattr(c_proj_mod, "weight", None)
                if o_w is not None:
                    o_w = o_w.data.clone()
                    if o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]:
                        o_w = o_w.t().contiguous()
                    o_b = getattr(c_proj_mod, "bias", None)
                    extracted["o"] = (
                        o_w,
                        o_b.data.clone() if o_b is not None else None,
                    )

        # Nested (BERT / RoBERTa)
        elif hasattr(pretrained_attn, "self") and hasattr(pretrained_attn.self, "query"):
            s = pretrained_attn.self
            for key, sub_name in (("q", "query"), ("k", "key"), ("v", "value")):
                m = _unwrap(getattr(s, sub_name, None))
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(pretrained_attn, "output") and hasattr(pretrained_attn.output, "dense"):
                out_mod = _unwrap(pretrained_attn.output.dense)
                b = getattr(out_mod, "bias", None)
                extracted["o"] = (
                    out_mod.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat (DistilBERT / generic)
        else:
            candidates = {
                "q": ("q_lin", "q_proj", "query", "W_q"),
                "k": ("k_lin", "k_proj", "key", "W_k"),
                "v": ("v_lin", "v_proj", "value", "W_v"),
                "o": ("out_lin", "out_proj", "dense", "o_proj"),
            }
            for key, names in candidates.items():
                for n in names:
                    if hasattr(pretrained_attn, n):
                        m = _unwrap(getattr(pretrained_attn, n))
                        if m is not None and hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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

                ad = UniversalAdapter()
                for k in missing:
                    w, b = ad.get_weight_bias(pretrained_attn, 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)

        # CS-FIX-18: log missing but do not raise
        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:
            logger.warning(
                "ChelaSlayAttention: missing projections %s from %s; "
                "using initialized weights for those.",
                missing_after,
                pretrained_attn.__class__.__name__,
            )

        target_map = {"q": self.W_q, "k": self.W_k, "v": self.W_v, "o": self.out_proj}
        with torch.no_grad():
            for key, target_mod in target_map.items():
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target_mod.weight.shape
                    and w.t().contiguous().shape == target_mod.weight.shape
                ):
                    w = w.t().contiguous()

                # CS-FIX-21: never zero the whole matrix; only copy overlap
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))
                if out_d == target_mod.weight.size(0) and in_d == target_mod.weight.size(1):
                    target_mod.weight.data.copy_(
                        w.to(target_mod.weight.device, target_mod.weight.dtype)
                    )
                else:
                    target_mod.weight.data.zero_()
                    target_mod.weight.data[:out_d, :in_d].copy_(
                        w[:out_d, :in_d].to(target_mod.weight.device, target_mod.weight.dtype)
                    )

                if b is not None and target_mod.bias is not None:
                    mn = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:mn].copy_(
                        b[:mn].to(target_mod.bias.device, target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                nrm = float(torch.norm(target_mod.weight.data).item())
                if nrm < 1e-4:
                    logger.warning("ChelaSlay: '%s' near-zero norm (%.2e)", key, nrm)

                if self.freeze_base:
                    target_mod.freeze_base_weights()

        try:
            for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

    # ------------------------------------------------------------------
    def _spherical_normalize(self, x: torch.Tensor) -> torch.Tensor:
        norm = torch.norm(x, p=2, dim=-1, keepdim=True)
        magnitude = torch.log1p(norm.clamp_min(0.0))
        return x / (norm + 1e-6) * magnitude

    def _yat_feature_map(self, x: torch.Tensor) -> torch.Tensor:
        if self.yat_proj.shape[0] != x.shape[-1]:
            raise RuntimeError(
                f"ChelaSlay: yat_proj shape {tuple(self.yat_proj.shape)} "
                f"vs x last dim {x.shape[-1]}."
            )
        proj = torch.matmul(x, self.yat_proj.to(x.device, x.dtype))
        scale = math.sqrt(2.0 / float(self.feature_dim))
        return torch.cat([torch.cos(proj), torch.sin(proj)], dim=-1) * scale

    # ------------------------------------------------------------------
    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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ):
        x = hidden_states
        B, L, D = x.shape
        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else x

        q = self.W_q(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k = (
            self.W_k(kv_input)
            .view(B, -1, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.W_v(kv_input)
            .view(B, -1, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        kv_len = k.shape[2]

        if position_embeddings is not None and not is_cross:
            if isinstance(position_embeddings, tuple):
                cos, sin = position_embeddings
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        elif self.rotary_emb is not None and not is_cross:
            if position_ids is not None:
                try:
                    cos, sin = self.rotary_emb(x, seq_len=kv_len)
                except TypeError:
                    cos, sin = self.rotary_emb(x, position_ids)
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)

        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        # Dense stream
        dense_scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * (
            self.head_dim**-0.5
        )

        can_causal = bool(self.is_causal and not is_cross and kv_len == L)
        if can_causal:
            causal = torch.triu(
                torch.full((L, L), -1e4, device=x.device, dtype=torch.float32),
                diagonal=1,
            )
            dense_scores = dense_scores + causal.unsqueeze(0).unsqueeze(0)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        valid_mask = None
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                valid_mask = (eff_mask > 0).float().unsqueeze(1).unsqueeze(-1)
                pad_m = (1.0 - eff_mask.unsqueeze(1).unsqueeze(2).float()) * -1e4
                dense_scores = dense_scores + pad_m
            elif eff_mask.dim() == 3:
                valid_mask = (eff_mask[:, 0, :] > -1000.0).float().unsqueeze(1).unsqueeze(-1)
                dense_scores = dense_scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                valid_mask = (eff_mask[:, 0, -1, :] > -1000.0).float().unsqueeze(1).unsqueeze(-1)
                dense_scores = dense_scores + eff_mask.float()

        dense_scores = torch.nan_to_num(dense_scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        dense_attn = F.softmax(dense_scores, dim=-1)
        dense_attn = torch.nan_to_num(dense_attn, nan=0.0)

        head_mask_b = None
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask_b = head_mask.view(1, -1, 1, 1)
            elif head_mask.dim() == 2:
                head_mask_b = head_mask.view(head_mask.shape[0], head_mask.shape[1], 1, 1)
            elif head_mask.dim() == 4:
                head_mask_b = head_mask
            if head_mask_b is not None:
                dense_attn = dense_attn * head_mask_b.to(dense_attn.dtype)

        dense_attn = self.dropout(dense_attn)
        dense_out = torch.matmul(dense_attn.to(v.dtype), v)

        # Linear stream
        q_norm = self._spherical_normalize(q)
        k_norm = self._spherical_normalize(k)
        q_feat = self._yat_feature_map(q_norm)
        k_feat = self._yat_feature_map(k_norm)

        if valid_mask is not None:
            k_feat = k_feat * valid_mask.to(k_feat.device, k_feat.dtype)
            v_masked = v * valid_mask.to(v.device, v.dtype)
        else:
            v_masked = v

        if can_causal:
            linear_out = self._causal_linear_attention(q_feat, k_feat, v_masked, L)
        else:
            linear_out = self._bidirectional_linear_attention(q_feat, k_feat, v_masked)

        if head_mask_b is not None:
            hw = head_mask_b.squeeze(-1).squeeze(-1)
            linear_out = linear_out * hw.unsqueeze(-1).to(linear_out.dtype)

        # Adaptive gating
        g = torch.sigmoid(self.lora_gate(x)).transpose(1, 2).contiguous().unsqueeze(-1)
        fused = g * dense_out + (1.0 - g) * linear_out
        out = fused.transpose(1, 2).contiguous().view(B, L, self.all_head_size)

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

        return final_out, (dense_attn if output_attentions else None)

    # ------------------------------------------------------------------
    def _bidirectional_linear_attention(self, q_feat, k_feat, v_masked):
        KV = torch.matmul(k_feat.transpose(-2, -1).contiguous().float(), v_masked.float())
        Z = k_feat.float().sum(dim=2, keepdim=True)
        num = torch.matmul(q_feat.float(), KV)
        denom = (q_feat.float() * Z).sum(dim=-1, keepdim=True)
        denom = denom.abs() + 1e-5
        return (num / denom).to(q_feat.dtype)

    def _causal_linear_attention(self, q_feat, k_feat, v_masked, L):
        chunk_size = 64
        B, H, _, feat_2f = k_feat.shape
        Dh = v_masked.shape[-1]
        num_chunks = (L + chunk_size - 1) // chunk_size

        S_running = torch.zeros(B, H, feat_2f, Dh, device=q_feat.device, dtype=torch.float32)
        Z_running = torch.zeros(B, H, feat_2f, device=q_feat.device, dtype=torch.float32)
        out_chunks = []

        for c in range(num_chunks):
            s = c * chunk_size
            e = min(s + chunk_size, L)
            q_c = q_feat[:, :, s:e, :].float()
            k_c = k_feat[:, :, s:e, :].float()
            v_c = v_masked[:, :, s:e, :].float()

            outer = k_c.unsqueeze(-1) * v_c.unsqueeze(-2)
            S_intra = torch.cumsum(outer, dim=2)
            Z_intra = torch.cumsum(k_c, dim=2)

            S_total = S_intra + S_running.unsqueeze(2)
            Z_total = Z_intra + Z_running.unsqueeze(2)

            num_c = torch.matmul(q_c.unsqueeze(-2), S_total).squeeze(-2)
            denom_c = (q_c * Z_total).sum(dim=-1, keepdim=True)
            denom_c = denom_c.abs() + 1e-5
            out_chunks.append((num_c / denom_c).to(q_feat.dtype))

            S_running = S_running + outer.sum(dim=2)
            Z_running = Z_running + k_c.sum(dim=2)

        return torch.cat(out_chunks, dim=2)

    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("W_q", self.W_q),
            ("W_k", self.W_k),
            ("W_v", self.W_v),
            ("out_proj", self.out_proj),
            ("lora_gate", self.lora_gate),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["yat_proj"] = _sz(self.yat_proj)
        out["total"] = sum(out.values())
        return out

CHELA-SLAY: Spherical Linearized Attention with Yat-Kernel.

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.

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 gate

Expand source code

@property
def gate(self):
    return self.lora_gate

prop k_lin

Expand source code

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

prop k_proj

Expand source code

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

prop key

Expand source code

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

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.W_q

prop q_proj

Expand source code

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

prop query

Expand source code

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

prop random_proj

Expand source code

@property
def random_proj(self):
    return self.yat_proj

prop v_lin

Expand source code

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

prop v_proj

Expand source code

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

prop value

Expand source code

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

Methods

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("W_q", self.W_q),
        ("W_k", self.W_k),
        ("W_v", self.W_v),
        ("out_proj", self.out_proj),
        ("lora_gate", self.lora_gate),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["yat_proj"] = _sz(self.yat_proj)
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class ClusterAttention (config: Any | None = None,
rank: int | None = None,
num_clusters: int = 16,
n_iters: int = 3,
window_size: int = 32,
lora_rank: int = 16,
comp_config: Any | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "cluster", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "cluster_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class ClusterAttention(BaseAttention):
    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,
        num_clusters: int = 16,
        n_iters: int = 3,
        window_size: int = 32,
        lora_rank: int = 16,
        comp_config: Any | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        # Resolve hidden_size and num_heads safely even if config is None
        hidden_size = kwargs.get("d_model") or kwargs.get("hidden_size") or kwargs.get("dim")
        num_heads = (
            kwargs.get("num_heads") or kwargs.get("n_heads") or kwargs.get("num_attention_heads")
        )
        num_kv_heads = kwargs.get("num_kv_heads") or kwargs.get("num_key_value_heads")

        if not num_heads and source_attn is not None:
            num_heads = getattr(
                source_attn,
                "num_heads",
                getattr(
                    source_attn,
                    "num_attention_heads",
                    getattr(source_attn, "n_heads", None),
                ),
            )
        if not num_heads and config is not None:
            num_heads = getattr(config, "num_attention_heads", getattr(config, "n_heads", None))
        num_heads = int(num_heads or 12)

        if not hidden_size and source_attn is not None:
            cand = getattr(
                source_attn,
                "hidden_size",
                getattr(
                    source_attn,
                    "embed_dim",
                    getattr(source_attn, "all_head_size", None),
                ),
            )
            if (
                cand is None
                and hasattr(source_attn, "query")
                and hasattr(source_attn.query, "in_features")
            ):
                cand = source_attn.query.in_features
            if (
                cand is None
                and hasattr(source_attn, "q_proj")
                and hasattr(source_attn.q_proj, "in_features")
            ):
                cand = source_attn.q_proj.in_features
            if (
                cand is None
                and hasattr(source_attn, "q_lin")
                and hasattr(source_attn.q_lin, "in_features")
            ):
                cand = source_attn.q_lin.in_features
            if (
                cand is None
                and hasattr(source_attn, "c_attn")
                and hasattr(source_attn.c_attn, "in_features")
            ):
                cand = source_attn.c_attn.in_features
            hidden_size = cand
        if not hidden_size and config is not None:
            hidden_size = getattr(config, "hidden_size", getattr(config, "dim", None))
        hidden_size = int(hidden_size or 768)

        head_dim = kwargs.get("head_dim", 0)
        if not head_dim and source_attn is not None:
            head_dim = getattr(source_attn, "head_dim", 0)
            if not head_dim:
                for attr in ("q_proj", "query", "q_lin"):
                    mod = getattr(source_attn, attr, None)
                    if mod is not None and hasattr(mod, "weight"):
                        w = mod.weight
                        head_dim = w.shape[0] // max(1, num_heads)
                        break

        if not head_dim and config is not None:
            head_dim = getattr(config, "head_dim", 0)
            if not head_dim:
                try:
                    adapter = ModelAdapter(config)
                    head_dim = getattr(adapter, "head_dim", 0)
                except Exception:
                    pass

        if not head_dim:
            head_dim = hidden_size // max(1, num_heads) if num_heads > 0 else 64
        head_dim = int(head_dim)

        if not num_kv_heads and source_attn is not None:
            for attr in ("k_proj", "key", "k_lin"):
                mod = getattr(source_attn, attr, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    num_kv_heads = w.shape[0] // max(1, head_dim)
                    break
            if not num_kv_heads:
                num_kv_heads = getattr(source_attn, "num_key_value_heads", None) or getattr(
                    source_attn, "num_kv_heads", None
                )

        if not num_kv_heads and config is not None:
            num_kv_heads = getattr(config, "num_key_value_heads", None)

        num_kv_heads = int(num_kv_heads or num_heads)

        self.hidden_size: int = hidden_size
        self.d_model: int = hidden_size
        self.num_heads: int = num_heads
        self.num_kv_heads: int = num_kv_heads
        self.head_dim: int = 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 = max(1, self.num_heads // max(1, self.num_kv_heads))

        if config is not None:
            self.num_clusters = getattr(config, "cluster_num_clusters", num_clusters)
        else:
            self.num_clusters = num_clusters
        # Fast K-means iterations (default 3, removed max(20, n_iters) constraint)
        self.n_iters: int = max(1, n_iters)
        self.window_size: int = window_size
        self.temp = nn.Parameter(torch.tensor(0.1))
        self.local_residual = nn.Parameter(torch.tensor(0.5))

        # Pre-allocated index buffer for efficient windowed attention
        self.register_buffer("_seq_idx", torch.arange(4096), persistent=False)

        # Decoder & nested encoder detection with robust prefix matching
        model_type = str(getattr(config, "model_type", "") or "").lower() if config else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_dec = (
            (
                getattr(config, "is_decoder", False)
                or getattr(config, "is_causal", False)
                or getattr(config, "causal", False)
                or any(model_type.startswith(p) for p in _decoder_prefixes)
            )
            if config
            else False
        )
        if hasattr(config, "architectures") and any(
            "GPT" in str(a) or "CausalLM" in str(a)
            for a in (getattr(config, "architectures", []) or [])
        ):
            is_dec = True
        if source_attn is not None:
            cls_name = source_attn.__class__.__name__.lower()
            if (
                any(p in cls_name for p in ("gpt", "causal", "qwen", "llama", "mistral"))
                or getattr(source_attn, "is_causal", False)
                or getattr(source_attn, "is_decoder", False)
                or hasattr(source_attn, "c_attn")
            ):
                is_dec = True
        if kwargs.get("is_decoder") is not None:
            is_dec = bool(kwargs["is_decoder"])
        elif kwargs.get("is_causal") is not None:
            is_dec = bool(kwargs["is_causal"])
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)
        self.is_nested_encoder: bool = bool(kwargs.get("is_nested_encoder", False))

        if config is not None:
            self.lora_rank = getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", lora_rank)),
            )
            self.lora_alpha = getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        elif comp_config is not None:
            self.lora_rank = getattr(comp_config, "lora_rank", kwargs.get("lora_rank", lora_rank))
            self.lora_alpha = getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
        else:
            self.lora_rank = kwargs.get("lora_rank", lora_rank)
            self.lora_alpha = kwargs.get("lora_alpha", 32.0)
        self.freeze_base: bool = self.lora_rank > 0

        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=self.freeze_base,
        )
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                ),
            )
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )

        if source_attn is not None:
            self.copy_weights(source_attn, **kwargs)

    def copy_weights(self, src: Any, **kwargs: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. HuggingFace / standard 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 as e:
                logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

        # 2. GPT-2 Fused c_attn Conv1D handling (including LoRALinear wrappers)
        if (
            hasattr(src, "c_attn")
            and src.c_attn is not None
            and (
                extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
            )
        ):
            c_attn_mod = src.c_attn
            if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
                try:
                    c_attn_mod.merge()
                except Exception:
                    pass
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                p_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (weight.ndim == 2 and weight.shape[1] == 3 * self.hidden_size)
                ):
                    weight = weight.t().contiguous()
                split_size = weight.shape[0] // 3
                bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias_mod.data.clone() if bias_mod is not None else None
                extracted["q"] = (
                    weight[:split_size, :],
                    bias_data[:split_size] if bias_data is not None else None,
                )
                extracted["k"] = (
                    weight[split_size : 2 * split_size, :],
                    (bias_data[split_size : 2 * split_size] if bias_data is not None else None),
                )
                extracted["v"] = (
                    weight[2 * split_size :, :],
                    bias_data[2 * split_size :] if bias_data is not None else None,
                )

        if hasattr(src, "c_proj") and src.c_proj is not None and extracted["o"][0] is None:
            c_proj_mod = src.c_proj
            if hasattr(c_proj_mod, "merge") and not getattr(c_proj_mod, "merged", False):
                try:
                    c_proj_mod.merge()
                except Exception:
                    pass
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_o, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (w_o.ndim == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                extracted["o"] = (
                    w_o,
                    bias_o.data.clone() if bias_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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
                or (
                    ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                    and not has_out_proj
                )
            )
        ):
            is_nested = True
        if hasattr(src, "config"):
            if (
                getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
                and not has_out_proj
            ):
                is_nested = True
        if kwargs.get("is_nested_encoder") is not None:
            is_nested = kwargs["is_nested_encoder"]
        self.is_nested_encoder = bool(is_nested)

        # Check required weights
        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:
            raise RuntimeError(
                f"ClusterAttention missing required weights for projections: {missing}"
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    raise RuntimeError(
                        f"ClusterAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                raise RuntimeError(
                    f"ClusterAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
                )

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

                target.freeze_base_weights()

        if hasattr(src, "rotary_emb") and src.rotary_emb is not None:
            self.rotary_emb = src.rotary_emb

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

    def _kmeans_cluster(
        self,
        Q: torch.Tensor,
        K_num: int,
        iters: int = 3,
        valid_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        r"""
        Vectorized Lloyd's $K$-means clustering for queries with padding token filtering.
        Q: [B, H, L, d], valid_mask: Optional[torch.Tensor] of shape [B, L] (True for real tokens)
        -> centroids: [B, H, K, d], cluster_ids: [B, H, L], counts: [B, H, K, 1]
        """
        B, H, L, d = Q.shape
        K = min(K_num, L)

        if valid_mask is not None:
            valid_f = valid_mask.float()  # [B, L]
            num_valid = valid_f.sum(dim=-1, keepdim=True).clamp(min=1.0)  # [B, 1]
            step = (num_valid / K).clamp(min=1.0)  # [B, 1]
            k_idx = torch.arange(K, device=Q.device, dtype=torch.float32).unsqueeze(0)  # [1, K]
            target_ranks = (k_idx * step).clamp(max=num_valid - 1.0).long()  # [B, K]

            # Assign rank to each valid token: 0, 1, ..., num_valid-1
            ranks = (valid_f.cumsum(dim=-1) - 1.0) * valid_f  # [B, L]
            is_valid_3d = valid_mask.unsqueeze(1)  # [B, 1, L]
            rank_diff = torch.where(
                is_valid_3d,
                (ranks.unsqueeze(1) - target_ranks.unsqueeze(-1)).abs(),
                torch.tensor(1e6, device=Q.device),
            )
            chosen_indices = torch.argmin(rank_diff, dim=-1)  # [B, K]

            idx_expanded = chosen_indices.unsqueeze(1).unsqueeze(-1).expand(B, H, K, d)
            centroids = torch.gather(Q, 2, idx_expanded).clone()
        else:
            step = max(1, L // K)
            init_indices = torch.arange(0, L, step, device=Q.device)[:K]
            if len(init_indices) < K:
                pad_idx = torch.zeros(K - len(init_indices), dtype=torch.long, device=Q.device)
                init_indices = torch.cat([init_indices, pad_idx])
            centroids = Q[:, :, init_indices, :].clone()

        cluster_ids = torch.zeros(B, H, L, dtype=torch.long, device=Q.device)
        counts = torch.ones(B, H, K, 1, dtype=Q.dtype, device=Q.device)
        Q_f = Q.float()

        for _ in range(iters):
            centroids_f = centroids.float()
            dist = torch.cdist(Q_f, centroids_f)
            if valid_mask is not None:
                # Disincentivize padding tokens from pulling centroids
                dist = dist.masked_fill(~valid_mask.unsqueeze(1).unsqueeze(-1), 1e6)
            cluster_ids = torch.argmin(dist, dim=-1)

            one_hot = F.one_hot(cluster_ids, num_classes=K).to(dtype=Q.dtype)
            if valid_mask is not None:
                one_hot = one_hot * valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype=Q.dtype)

            counts = one_hot.sum(dim=2).unsqueeze(-1)
            counts_clamped = counts.clamp(min=1.0)
            sum_q = torch.einsum("bhlk,bhld->bhkd", one_hot, Q)
            centroids = torch.where(counts > 0, sum_q / (counts_clamped + 1e-8), centroids)
            centroids = torch.nan_to_num(centroids, nan=0.0)

        return centroids, cluster_ids, counts

    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,
        *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, L, D = hidden_states.shape
        N_k = kv_input.shape[1]
        H, Dh = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = self.q_proj(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, N_k, self.num_kv_heads, Dh).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, N_k, self.num_kv_heads, Dh).transpose(1, 2).contiguous()

        # Apply RoPE (accept from position_embeddings or fallback to rotary_emb)
        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin, kwargs.get("position_ids"))
        elif self.rotary_emb is not None:
            rotary_mod = self.rotary_emb
            pos_ids = kwargs.get("position_ids")
            if pos_ids is None:
                pos_ids = torch.arange(L, device=device).unsqueeze(0).expand(B, -1)
            try:
                cos, sin = rotary_mod(Q, seq_len=L)
            except Exception:
                try:
                    cos, sin = rotary_mod(Q, pos_ids)
                except Exception:
                    try:
                        cos, sin = rotary_mod(Q)
                    except Exception:
                        cos, sin = None, None
            if cos is not None and sin is not None:
                Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids=pos_ids)

        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = Dh**-0.5

        # Check for causal decoding (guard against self.config being None)
        _cfg = self.config
        is_dec_cfg = False
        if _cfg is not None:
            model_type = str(getattr(_cfg, "model_type", "") or "").lower()
            _decoder_prefixes = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            is_dec_cfg = (
                getattr(_cfg, "is_decoder", False)
                or getattr(_cfg, "is_causal", False)
                or getattr(_cfg, "causal", False)
                or any(model_type.startswith(p) for p in _decoder_prefixes)
            )
        is_decoder = self.is_decoder or is_dec_cfg

        if is_decoder or is_cross_attention:
            # Exact attention with causal masking to prevent future token leakage
            scores = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale
            mask = attention_mask if attention_mask is not None else kwargs.get("mask")
            if mask is not None:
                if mask.dim() == 4:
                    if mask.dtype == torch.bool:
                        scores = scores.masked_fill(~mask, -10000.0)
                    else:
                        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)

            if is_decoder:
                causal_tril = (
                    torch.tril(torch.ones((L, N_k), device=device, dtype=torch.bool))
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores.masked_fill(~causal_tril, -10000.0)

            attn = torch.nan_to_num(F.softmax(scores.float(), dim=-1), nan=0.0).to(dtype)
            attn = self.dropout(attn)

            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)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(-1)
                attn = attn * head_mask

            out_context = torch.matmul(attn, V)
            if not is_decoder:
                valid_q = get_token_valid_mask(attention_mask, B, L, device)
                if valid_q is not None:
                    out_context = out_context * valid_q.unsqueeze(1).unsqueeze(-1).to(dtype)
            attn_ret = attn if output_attentions else None
        else:
            # Bidirectional clustering attention
            valid_mask_2d = get_token_valid_mask(attention_mask, B, L, device)
            key_valid_4d = (
                valid_mask_2d.unsqueeze(1).unsqueeze(2)
                if valid_mask_2d is not None
                else torch.ones((B, 1, 1, L), dtype=torch.bool, device=device)
            )

            if L <= 64 and (self.num_clusters >= 8 or not output_attentions):
                scores = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale
                scores = scores.masked_fill(~key_valid_4d, -10000.0)
                attn = torch.nan_to_num(F.softmax(scores.float(), dim=-1), nan=0.0).to(dtype)
                attn = self.dropout(attn)
                out_context = torch.matmul(attn, V)
                attn_ret = attn if output_attentions else None
            else:
                K_num = min(max(1, self.num_clusters), L)

                # 1. Cluster Queries into K centroids with padding mask filtering
                centroids, cluster_ids, counts = self._kmeans_cluster(
                    Q, K_num, self.n_iters, valid_mask=valid_mask_2d
                )

                # 2. Centroid-to-Key Global Cross-Attention [B, H, K, L]
                scores_cent = torch.matmul(centroids, K.transpose(-2, -1).contiguous()) * scale
                scores_cent = scores_cent.masked_fill(~key_valid_4d, -10000.0)
                attn_cent = torch.nan_to_num(F.softmax(scores_cent.float(), dim=-1), nan=0.0).to(
                    dtype
                )
                attn_cent = self.dropout(attn_cent)

                # Centroid values: [B, H, K, Dh]
                V_cent = torch.matmul(attn_cent, V)

                # 3. Differentiable Query-to-Centroid soft assignment with learned sharpening & context gathering [B, H, L, Dh]
                scores_assign = torch.matmul(Q, centroids.transpose(-2, -1).contiguous()) * scale
                empty_clusters = (counts == 0).squeeze(-1).unsqueeze(2)  # [B, H, 1, K]
                scores_assign = scores_assign.masked_fill(empty_clusters, -10000.0)
                temp_clamped = self.temp.abs().clamp(min=0.1, max=10.0)
                attn_assign = torch.nan_to_num(
                    F.softmax((scores_assign / temp_clamped).float(), dim=-1), nan=0.0
                ).to(dtype)
                assign_denom = attn_assign.sum(dim=-1, keepdim=True).clamp(min=1.0)
                attn_assign = attn_assign / assign_denom
                global_context = torch.matmul(attn_assign, V_cent)

                # 4. Optional Local Sliding Window Attention
                if self.window_size > 0 and self.window_size < L:
                    scores_loc = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale
                    idx = (
                        self._seq_idx[:L].to(device)
                        if hasattr(self, "_seq_idx") and len(self._seq_idx) >= L
                        else torch.arange(L, device=device)
                    )
                    band_mask = (idx.unsqueeze(1) - idx.unsqueeze(0)).abs() <= self.window_size
                    band_mask_4d = band_mask.unsqueeze(0).unsqueeze(0) & key_valid_4d
                    scores_loc = scores_loc.masked_fill(~band_mask_4d, -10000.0)
                    attn_loc = torch.nan_to_num(F.softmax(scores_loc.float(), dim=-1), nan=0.0).to(
                        dtype
                    )
                    attn_loc = self.dropout(attn_loc)
                    local_context = torch.matmul(attn_loc, V)
                    alpha = torch.sigmoid(self.local_residual)
                    out_context = (1.0 - alpha) * global_context + alpha * local_context
                else:
                    out_context = global_context

                attn_ret = attn_cent if output_attentions else None

            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)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(-1)
                out_context = out_context * head_mask

            if valid_mask_2d is not None:
                out_context = out_context * valid_mask_2d.unsqueeze(1).unsqueeze(-1).to(dtype)

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

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

        valid_out = get_token_valid_mask(attention_mask, B, L, device)
        if valid_out is not None:
            output = output * valid_out.unsqueeze(-1).to(output.dtype)

        return output, attn_ret

Every 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_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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class CompressedAttention (config: Any | None = None,
rank: int | None = None,
rank_qkv: int = 256,
rank_out: int = 256,
lora_rank: int = 16,
original_layer: Any = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "compressed_attention_layer",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "compressed_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class CompressedAttentionLayer(BaseAttention):
    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,
        rank_qkv: int = 256,
        rank_out: int = 256,
        lora_rank: int = 16,
        original_layer: Any = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config if config is not None else original_layer
        super().__init__(cfg, rank)
        self.adapter = ModelAdapter(cfg)
        self.hidden_size: int = self.adapter.hidden_size or kwargs.get(
            "hidden_size", kwargs.get("d_model", 768)
        )
        self.num_heads: int = self.adapter.num_heads or kwargs.get(
            "num_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads: int = self.adapter.num_key_value_heads or kwargs.get(
            "num_kv_heads", kwargs.get("num_key_value_heads", self.num_heads)
        )
        self.head_dim: int = self.adapter.head_dim or (self.hidden_size // max(1, self.num_heads))
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))
        _mt = getattr(cfg, "model_type", "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        self.is_decoder: bool = bool(
            getattr(cfg, "is_decoder", False) or any(_mt.startswith(p) for p in _decoder_prefixes)
        )
        self.is_nested_encoder: bool = False

        rank_val = rank or kwargs.get("rank_qkv", rank_qkv)
        self.rank_qkv = min(rank_val, self.all_head_size + 2 * self.kv_head_size, self.hidden_size)
        self.rank_out = min(
            rank or kwargs.get("rank_out", rank_out),
            self.all_head_size,
            self.hidden_size,
        )
        self.lora_rank = getattr(cfg, "lora_rank", lora_rank)
        self.lora_scaling = (32.0 / self.lora_rank) if self.lora_rank > 0 else 1.0

        qkv_out_dim = self.all_head_size + 2 * self.kv_head_size

        # Frozen base SVD parameters
        self.Vh_qkv = nn.Parameter(
            torch.randn(self.rank_qkv, self.hidden_size) * 0.02, requires_grad=False
        )
        self.S_qkv_sqrt = nn.Parameter(torch.ones(self.rank_qkv), requires_grad=False)
        self.U_qkv = nn.Parameter(
            torch.randn(qkv_out_dim, self.rank_qkv) * 0.02, requires_grad=False
        )
        self.bias_qkv = nn.Parameter(torch.zeros(qkv_out_dim), requires_grad=False)

        self.Vh_out = nn.Parameter(
            torch.randn(self.rank_out, self.all_head_size) * 0.02, requires_grad=False
        )
        self.S_out_sqrt = nn.Parameter(torch.ones(self.rank_out), requires_grad=False)
        self.U_out = nn.Parameter(
            torch.randn(self.hidden_size, self.rank_out) * 0.02, requires_grad=False
        )
        self.bias_out = nn.Parameter(torch.zeros(self.hidden_size), requires_grad=False)

        # Trainable LoRA residual adapters
        if self.lora_rank > 0:
            self.lora_A_qkv = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.lora_B_qkv = nn.Parameter(torch.empty(qkv_out_dim, self.lora_rank))
            self.lora_A_out = nn.Parameter(torch.empty(self.lora_rank, self.all_head_size))
            self.lora_B_out = nn.Parameter(torch.empty(self.hidden_size, self.lora_rank))
            self.reset_lora_parameters()
        else:
            self.register_parameter("lora_A_qkv", None)
            self.register_parameter("lora_B_qkv", None)
            self.register_parameter("lora_A_out", None)
            self.register_parameter("lora_B_out", None)

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", 0.1))
            if cfg
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

        src = pretrained_attn or (
            original_layer
            if hasattr(original_layer, "q_lin")
            or hasattr(original_layer, "query")
            or hasattr(original_layer, "c_attn")
            else None
        )
        if src is not None:
            self.copy_weights(src)

    @property
    def qkv_weight(self) -> torch.Tensor:
        """Reconstruct the effective base QKV weight from SVD factors."""
        return (self.U_qkv * self.S_qkv_sqrt.view(1, -1)) @ (
            self.Vh_qkv * self.S_qkv_sqrt.view(-1, 1)
        )

    @property
    def out_weight(self) -> torch.Tensor:
        """Reconstruct the effective base Out weight from SVD factors."""
        return (self.U_out * self.S_out_sqrt.view(1, -1)) @ (
            self.Vh_out * self.S_out_sqrt.view(-1, 1)
        )

    def reset_lora_parameters(self) -> None:
        if self.lora_rank > 0:
            nn.init.kaiming_uniform_(self.lora_A_qkv, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_qkv)
            nn.init.kaiming_uniform_(self.lora_A_out, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_out)
            self.lora_A_qkv.requires_grad = True
            self.lora_B_qkv.requires_grad = True
            self.lora_A_out.requires_grad = True
            self.lora_B_out.requires_grad = True

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. HuggingFace / standard 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 as e:
                logger.debug(
                    f"UniversalAdapter failed for {proj}: {e}"
                )  # continue to fallback paths

        # 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
        ):
            weight = src.c_attn.weight.data
            bias = getattr(src.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                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:
            w_o = src.c_proj.weight.data
            if (
                w_o.dim() == 2
                and w_o.shape[0] != self.hidden_size
                and w_o.shape[1] == self.hidden_size
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(src.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        # Check output projection attrs FIRST — DistilBERT has out_lin and is NOT nested
        if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
            self.is_nested_encoder = False
        elif (
            (
                hasattr(src, "self")
                and hasattr(src.self, "query")
                and hasattr(src, "output")
                and hasattr(src.output, "dense")
            )
            or (
                cls_name in ("BertSelfAttention", "RobertaSelfAttention")
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        # Check required weights
        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(
                f"CompressedAttentionLayer: missing weights for projections {missing}; proceeding with random init"
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"CompressedAttentionLayer: extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f}); continuing"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"CompressedAttentionLayer: extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f}); continuing"
                )

        # Joint SVD on Q, K, V
        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():
            w_q = _orient(extracted["q"][0].float(), self.hidden_size)
            w_k = _orient(extracted["k"][0].float(), self.hidden_size)
            w_v = _orient(extracted["v"][0].float(), self.hidden_size)
            w_qkv = torch.cat([w_q, w_k, w_v], dim=0)

            U, S, Vh = torch.linalg.svd(w_qkv, full_matrices=False)
            r_qkv = min(self.rank_qkv, U.size(1), Vh.size(0))
            sqrt_s = torch.sqrt(torch.clamp(S[:r_qkv], min=1e-12))

            qkv_out_dim = self.all_head_size + 2 * self.kv_head_size
            self.U_qkv.data[: min(qkv_out_dim, U.size(0)), :r_qkv].copy_(
                U[: min(qkv_out_dim, U.size(0)), :r_qkv]
            )
            self.S_qkv_sqrt.data[:r_qkv].copy_(sqrt_s)
            self.Vh_qkv.data[:r_qkv, : min(self.hidden_size, Vh.size(1))].copy_(
                Vh[:r_qkv, : min(self.hidden_size, Vh.size(1))]
            )

            b_q = (
                extracted["q"][1]
                if extracted["q"][1] is not None
                else torch.zeros(w_q.size(0), device=w_q.device)
            )
            b_k = (
                extracted["k"][1]
                if extracted["k"][1] is not None
                else torch.zeros(w_k.size(0), device=w_k.device)
            )
            b_v = (
                extracted["v"][1]
                if extracted["v"][1] is not None
                else torch.zeros(w_v.size(0), device=w_v.device)
            )
            b_q_val = b_q.data if isinstance(b_q, nn.Parameter) else b_q
            b_k_val = b_k.data if isinstance(b_k, nn.Parameter) else b_k
            b_v_val = b_v.data if isinstance(b_v, nn.Parameter) else b_v
            b_qkv = torch.cat([b_q_val, b_k_val, b_v_val], dim=0)
            self.bias_qkv.data[: min(qkv_out_dim, b_qkv.size(0))].copy_(
                b_qkv[: min(qkv_out_dim, b_qkv.size(0))]
            )

            # Output projection SVD
            if extracted["o"][0] is not None:
                w_o = _orient(extracted["o"][0].float(), self.all_head_size)
                U_o, S_o, Vh_o = torch.linalg.svd(w_o, full_matrices=False)
                r_out = min(self.rank_out, U_o.size(1), Vh_o.size(0))
                sqrt_so = torch.sqrt(torch.clamp(S_o[:r_out], min=1e-12))

                self.U_out.data[: min(self.hidden_size, U_o.size(0)), :r_out].copy_(
                    U_o[: min(self.hidden_size, U_o.size(0)), :r_out]
                )
                self.S_out_sqrt.data[:r_out].copy_(sqrt_so)
                self.Vh_out.data[:r_out, : min(self.all_head_size, Vh_o.size(1))].copy_(
                    Vh_o[:r_out, : min(self.all_head_size, Vh_o.size(1))]
                )

                if extracted["o"][1] is not None:
                    b_o = extracted["o"][1]
                    b_o_data = b_o.data if isinstance(b_o, nn.Parameter) else b_o
                    self.bias_out.data[: min(self.hidden_size, b_o_data.size(0))].copy_(
                        b_o_data[: min(self.hidden_size, b_o_data.size(0))]
                    )
            elif self.is_nested_encoder:
                self.U_out.data.zero_()
                min_d = min(self.hidden_size, self.rank_out)
                self.U_out.data[:min_d, :min_d].fill_diagonal_(1.0)
                self.S_out_sqrt.data.fill_(1.0)
                self.Vh_out.data.zero_()
                self.Vh_out.data[:min_d, :min_d].fill_diagonal_(1.0)
                self.bias_out.data.zero_()

            # Ensure base weights are frozen
            self.Vh_qkv.requires_grad = False
            self.S_qkv_sqrt.requires_grad = False
            self.U_qkv.requires_grad = False
            self.bias_qkv.requires_grad = False
            self.Vh_out.requires_grad = False
            self.S_out_sqrt.requires_grad = False
            self.U_out.requires_grad = False
            self.bias_out.requires_grad = False

            if self.lora_rank > 0:
                self.reset_lora_parameters()

        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]:
        if hidden_states.dim() == 2:
            hidden_states = hidden_states.unsqueeze(0)
        B, N, D = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        # FIX 1: Pre-transpose weight matrices to avoid runtime transposes
        # Vh_qkv is (Rank, D), U_qkv is (Out_Dim, Rank)
        # S_qkv_sqrt is (Rank)

        # Optimize Matmul path
        v_proj = torch.matmul(hidden_states, self.Vh_qkv.t()) * self.S_qkv_sqrt
        qkv = torch.matmul(v_proj, (self.U_qkv * self.S_qkv_sqrt).t()) + self.bias_qkv

        # LoRA residual adapter
        if self.lora_rank > 0 and self.lora_A_qkv is not None and self.lora_B_qkv is not None:
            lora_qkv = (
                F.linear(F.linear(hidden_states, self.lora_A_qkv), self.lora_B_qkv)
                * self.lora_scaling
            )
            qkv = qkv + lora_qkv

        # Split into Q, K, V
        q, k, v = torch.split(
            qkv, [self.all_head_size, self.kv_head_size, self.kv_head_size], dim=-1
        )

        # FIX 2: Remove .contiguous() memory bombs
        q = q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # Apply RoPE
        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"))

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

        scale = 1.0 / math.sqrt(self.head_dim)

        # Use torch.baddbmm or scaled_dot_product_attention for huge memory savings
        # Fallback to standard matmul if not available or unsupported shapes
        try:
            # PyTorch 2.0+ Flash Attention / Memory Efficient Attention
            # This handles Causal masking automatically and saves massive VRAM
            is_causal = self.is_decoder and (attention_mask is None or attention_mask.dim() != 4)
            attn_out = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=attention_mask,
                dropout_p=self.dropout.p if self.training else 0.0,
                is_causal=is_causal,
            )
            probs = None  # SDPA doesn't return probabilities to save memory
        except Exception:
            # Legacy Path (if SDPA fails due to old PyTorch or weird masks)
            scores = torch.matmul(q, k.transpose(-2, -1)) * scale

            mask = attention_mask if attention_mask is not None else kwargs.get("mask")
            if mask is not None:
                if mask.dim() == 4:
                    if mask.dtype == torch.bool:
                        scores = scores.masked_fill(~mask, -10000.0)
                    else:
                        scores = scores + mask.to(scores.dtype)
                else:
                    valid_m = get_token_valid_mask(mask, B, N, device)
                    if valid_m is not None:
                        scores = scores.masked_fill((~valid_m).unsqueeze(1).unsqueeze(2), -10000.0)

            # Explicit Causal Masking
            if self.is_decoder and (mask is None or mask.dim() != 4):
                causal_tril = (
                    torch.tril(torch.ones((N, N), device=device, dtype=torch.bool))
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores.masked_fill(~causal_tril, -10000.0)

            # FIX 3: Safe Softmax without extra type casts that create new tensors
            probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(dtype)
            probs = self.dropout(probs)
            attn_out = torch.matmul(probs, v)

        # Transpose back and view (reshape handles non-contiguous memory safely without deep copies)
        attn_out = attn_out.transpose(1, 2).reshape(B, N, self.all_head_size)

        if self.is_nested_encoder:
            out = attn_out
        else:
            # SVD Out forward
            out_v = torch.matmul(attn_out, self.Vh_out.t()) * self.S_out_sqrt
            out = torch.matmul(out_v, (self.U_out * self.S_out_sqrt).t()) + self.bias_out
            if self.lora_rank > 0 and self.lora_A_out is not None and self.lora_B_out is not None:
                lora_out = (
                    F.linear(F.linear(attn_out, self.lora_A_out), self.lora_B_out)
                    * self.lora_scaling
                )
                out = out + lora_out

        return out, (probs if output_attentions and probs is not None else None)

Every 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

Subclasses

• CompressedAttentionLayer2

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 out_weight : torch.Tensor

Expand source code

@property
def out_weight(self) -> torch.Tensor:
    """Reconstruct the effective base Out weight from SVD factors."""
    return (self.U_out * self.S_out_sqrt.view(1, -1)) @ (
        self.Vh_out * self.S_out_sqrt.view(-1, 1)
    )

Reconstruct the effective base Out weight from SVD factors.

prop qkv_weight : torch.Tensor

Expand source code

@property
def qkv_weight(self) -> torch.Tensor:
    """Reconstruct the effective base QKV weight from SVD factors."""
    return (self.U_qkv * self.S_qkv_sqrt.view(1, -1)) @ (
        self.Vh_qkv * self.S_qkv_sqrt.view(-1, 1)
    )

Reconstruct the effective base QKV weight from SVD factors.

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    from neuralops.core.adapters import UniversalAdapter

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

    # 1. HuggingFace / standard 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 as e:
            logger.debug(
                f"UniversalAdapter failed for {proj}: {e}"
            )  # continue to fallback paths

    # 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
    ):
        weight = src.c_attn.weight.data
        bias = getattr(src.c_attn, "bias", None)
        bias_data = bias.data if bias is not None else None
        if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
            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:
        w_o = src.c_proj.weight.data
        if (
            w_o.dim() == 2
            and w_o.shape[0] != self.hidden_size
            and w_o.shape[1] == self.hidden_size
        ):
            w_o = w_o.t().contiguous()
        b_o = getattr(src.c_proj, "bias", None)
        extracted["o"] = (w_o, b_o.data if b_o is not None else None)

    # 3. RoBERTa / BERT Nested Encoder handling
    if hasattr(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"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    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 = src.__class__.__name__
    has_out_proj = any(
        hasattr(src, attr)
        for attr in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    # Check output projection attrs FIRST — DistilBERT has out_lin and is NOT nested
    if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
        self.is_nested_encoder = False
    elif (
        (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            and hasattr(src, "output")
            and hasattr(src.output, "dense")
        )
        or (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        )
        and not has_out_proj
    ):
        self.is_nested_encoder = True
    else:
        self.is_nested_encoder = False

    # Check required weights
    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(
            f"CompressedAttentionLayer: missing weights for projections {missing}; proceeding with random init"
        )

    # Frobenius norm validation
    for proj_key in ["q", "k", "v"]:
        w = extracted[proj_key][0]
        if w is not None:
            norm = torch.linalg.norm(w.float())
            if norm < 1e-4:
                logger.warning(
                    f"CompressedAttentionLayer: extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f}); continuing"
                )

    if extracted["o"][0] is not None:
        norm_o = torch.linalg.norm(extracted["o"][0].float())
        if norm_o < 1e-4:
            logger.warning(
                f"CompressedAttentionLayer: extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f}); continuing"
            )

    # Joint SVD on Q, K, V
    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():
        w_q = _orient(extracted["q"][0].float(), self.hidden_size)
        w_k = _orient(extracted["k"][0].float(), self.hidden_size)
        w_v = _orient(extracted["v"][0].float(), self.hidden_size)
        w_qkv = torch.cat([w_q, w_k, w_v], dim=0)

        U, S, Vh = torch.linalg.svd(w_qkv, full_matrices=False)
        r_qkv = min(self.rank_qkv, U.size(1), Vh.size(0))
        sqrt_s = torch.sqrt(torch.clamp(S[:r_qkv], min=1e-12))

        qkv_out_dim = self.all_head_size + 2 * self.kv_head_size
        self.U_qkv.data[: min(qkv_out_dim, U.size(0)), :r_qkv].copy_(
            U[: min(qkv_out_dim, U.size(0)), :r_qkv]
        )
        self.S_qkv_sqrt.data[:r_qkv].copy_(sqrt_s)
        self.Vh_qkv.data[:r_qkv, : min(self.hidden_size, Vh.size(1))].copy_(
            Vh[:r_qkv, : min(self.hidden_size, Vh.size(1))]
        )

        b_q = (
            extracted["q"][1]
            if extracted["q"][1] is not None
            else torch.zeros(w_q.size(0), device=w_q.device)
        )
        b_k = (
            extracted["k"][1]
            if extracted["k"][1] is not None
            else torch.zeros(w_k.size(0), device=w_k.device)
        )
        b_v = (
            extracted["v"][1]
            if extracted["v"][1] is not None
            else torch.zeros(w_v.size(0), device=w_v.device)
        )
        b_q_val = b_q.data if isinstance(b_q, nn.Parameter) else b_q
        b_k_val = b_k.data if isinstance(b_k, nn.Parameter) else b_k
        b_v_val = b_v.data if isinstance(b_v, nn.Parameter) else b_v
        b_qkv = torch.cat([b_q_val, b_k_val, b_v_val], dim=0)
        self.bias_qkv.data[: min(qkv_out_dim, b_qkv.size(0))].copy_(
            b_qkv[: min(qkv_out_dim, b_qkv.size(0))]
        )

        # Output projection SVD
        if extracted["o"][0] is not None:
            w_o = _orient(extracted["o"][0].float(), self.all_head_size)
            U_o, S_o, Vh_o = torch.linalg.svd(w_o, full_matrices=False)
            r_out = min(self.rank_out, U_o.size(1), Vh_o.size(0))
            sqrt_so = torch.sqrt(torch.clamp(S_o[:r_out], min=1e-12))

            self.U_out.data[: min(self.hidden_size, U_o.size(0)), :r_out].copy_(
                U_o[: min(self.hidden_size, U_o.size(0)), :r_out]
            )
            self.S_out_sqrt.data[:r_out].copy_(sqrt_so)
            self.Vh_out.data[:r_out, : min(self.all_head_size, Vh_o.size(1))].copy_(
                Vh_o[:r_out, : min(self.all_head_size, Vh_o.size(1))]
            )

            if extracted["o"][1] is not None:
                b_o = extracted["o"][1]
                b_o_data = b_o.data if isinstance(b_o, nn.Parameter) else b_o
                self.bias_out.data[: min(self.hidden_size, b_o_data.size(0))].copy_(
                    b_o_data[: min(self.hidden_size, b_o_data.size(0))]
                )
        elif self.is_nested_encoder:
            self.U_out.data.zero_()
            min_d = min(self.hidden_size, self.rank_out)
            self.U_out.data[:min_d, :min_d].fill_diagonal_(1.0)
            self.S_out_sqrt.data.fill_(1.0)
            self.Vh_out.data.zero_()
            self.Vh_out.data[:min_d, :min_d].fill_diagonal_(1.0)
            self.bias_out.data.zero_()

        # Ensure base weights are frozen
        self.Vh_qkv.requires_grad = False
        self.S_qkv_sqrt.requires_grad = False
        self.U_qkv.requires_grad = False
        self.bias_qkv.requires_grad = False
        self.Vh_out.requires_grad = False
        self.S_out_sqrt.requires_grad = False
        self.U_out.requires_grad = False
        self.bias_out.requires_grad = False

        if self.lora_rank > 0:
            self.reset_lora_parameters()

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

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    if hidden_states.dim() == 2:
        hidden_states = hidden_states.unsqueeze(0)
    B, N, D = hidden_states.shape
    device = hidden_states.device
    dtype = hidden_states.dtype

    # FIX 1: Pre-transpose weight matrices to avoid runtime transposes
    # Vh_qkv is (Rank, D), U_qkv is (Out_Dim, Rank)
    # S_qkv_sqrt is (Rank)

    # Optimize Matmul path
    v_proj = torch.matmul(hidden_states, self.Vh_qkv.t()) * self.S_qkv_sqrt
    qkv = torch.matmul(v_proj, (self.U_qkv * self.S_qkv_sqrt).t()) + self.bias_qkv

    # LoRA residual adapter
    if self.lora_rank > 0 and self.lora_A_qkv is not None and self.lora_B_qkv is not None:
        lora_qkv = (
            F.linear(F.linear(hidden_states, self.lora_A_qkv), self.lora_B_qkv)
            * self.lora_scaling
        )
        qkv = qkv + lora_qkv

    # Split into Q, K, V
    q, k, v = torch.split(
        qkv, [self.all_head_size, self.kv_head_size, self.kv_head_size], dim=-1
    )

    # FIX 2: Remove .contiguous() memory bombs
    q = q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
    k = k.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2)
    v = v.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2)

    # Apply RoPE
    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"))

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

    scale = 1.0 / math.sqrt(self.head_dim)

    # Use torch.baddbmm or scaled_dot_product_attention for huge memory savings
    # Fallback to standard matmul if not available or unsupported shapes
    try:
        # PyTorch 2.0+ Flash Attention / Memory Efficient Attention
        # This handles Causal masking automatically and saves massive VRAM
        is_causal = self.is_decoder and (attention_mask is None or attention_mask.dim() != 4)
        attn_out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attention_mask,
            dropout_p=self.dropout.p if self.training else 0.0,
            is_causal=is_causal,
        )
        probs = None  # SDPA doesn't return probabilities to save memory
    except Exception:
        # Legacy Path (if SDPA fails due to old PyTorch or weird masks)
        scores = torch.matmul(q, k.transpose(-2, -1)) * scale

        mask = attention_mask if attention_mask is not None else kwargs.get("mask")
        if mask is not None:
            if mask.dim() == 4:
                if mask.dtype == torch.bool:
                    scores = scores.masked_fill(~mask, -10000.0)
                else:
                    scores = scores + mask.to(scores.dtype)
            else:
                valid_m = get_token_valid_mask(mask, B, N, device)
                if valid_m is not None:
                    scores = scores.masked_fill((~valid_m).unsqueeze(1).unsqueeze(2), -10000.0)

        # Explicit Causal Masking
        if self.is_decoder and (mask is None or mask.dim() != 4):
            causal_tril = (
                torch.tril(torch.ones((N, N), device=device, dtype=torch.bool))
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores.masked_fill(~causal_tril, -10000.0)

        # FIX 3: Safe Softmax without extra type casts that create new tensors
        probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(dtype)
        probs = self.dropout(probs)
        attn_out = torch.matmul(probs, v)

    # Transpose back and view (reshape handles non-contiguous memory safely without deep copies)
    attn_out = attn_out.transpose(1, 2).reshape(B, N, self.all_head_size)

    if self.is_nested_encoder:
        out = attn_out
    else:
        # SVD Out forward
        out_v = torch.matmul(attn_out, self.Vh_out.t()) * self.S_out_sqrt
        out = torch.matmul(out_v, (self.U_out * self.S_out_sqrt).t()) + self.bias_out
        if self.lora_rank > 0 and self.lora_A_out is not None and self.lora_B_out is not None:
            lora_out = (
                F.linear(F.linear(attn_out, self.lora_A_out), self.lora_B_out)
                * self.lora_scaling
            )
            out = out + lora_out

    return out, (probs if output_attentions and probs is not None else None)

Must return (output, attention_weights_or_None).

def reset_lora_parameters(self) ‑> None

Expand source code

def reset_lora_parameters(self) -> None:
    if self.lora_rank > 0:
        nn.init.kaiming_uniform_(self.lora_A_qkv, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B_qkv)
        nn.init.kaiming_uniform_(self.lora_A_out, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B_out)
        self.lora_A_qkv.requires_grad = True
        self.lora_B_qkv.requires_grad = True
        self.lora_A_out.requires_grad = True
        self.lora_B_out.requires_grad = True

class CompressedAttentionLayer (config: Any | None = None,
rank: int | None = None,
rank_qkv: int = 256,
rank_out: int = 256,
lora_rank: int = 16,
original_layer: Any = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "compressed_attention_layer",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "compressed_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class CompressedAttentionLayer(BaseAttention):
    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,
        rank_qkv: int = 256,
        rank_out: int = 256,
        lora_rank: int = 16,
        original_layer: Any = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config if config is not None else original_layer
        super().__init__(cfg, rank)
        self.adapter = ModelAdapter(cfg)
        self.hidden_size: int = self.adapter.hidden_size or kwargs.get(
            "hidden_size", kwargs.get("d_model", 768)
        )
        self.num_heads: int = self.adapter.num_heads or kwargs.get(
            "num_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads: int = self.adapter.num_key_value_heads or kwargs.get(
            "num_kv_heads", kwargs.get("num_key_value_heads", self.num_heads)
        )
        self.head_dim: int = self.adapter.head_dim or (self.hidden_size // max(1, self.num_heads))
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))
        _mt = getattr(cfg, "model_type", "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        self.is_decoder: bool = bool(
            getattr(cfg, "is_decoder", False) or any(_mt.startswith(p) for p in _decoder_prefixes)
        )
        self.is_nested_encoder: bool = False

        rank_val = rank or kwargs.get("rank_qkv", rank_qkv)
        self.rank_qkv = min(rank_val, self.all_head_size + 2 * self.kv_head_size, self.hidden_size)
        self.rank_out = min(
            rank or kwargs.get("rank_out", rank_out),
            self.all_head_size,
            self.hidden_size,
        )
        self.lora_rank = getattr(cfg, "lora_rank", lora_rank)
        self.lora_scaling = (32.0 / self.lora_rank) if self.lora_rank > 0 else 1.0

        qkv_out_dim = self.all_head_size + 2 * self.kv_head_size

        # Frozen base SVD parameters
        self.Vh_qkv = nn.Parameter(
            torch.randn(self.rank_qkv, self.hidden_size) * 0.02, requires_grad=False
        )
        self.S_qkv_sqrt = nn.Parameter(torch.ones(self.rank_qkv), requires_grad=False)
        self.U_qkv = nn.Parameter(
            torch.randn(qkv_out_dim, self.rank_qkv) * 0.02, requires_grad=False
        )
        self.bias_qkv = nn.Parameter(torch.zeros(qkv_out_dim), requires_grad=False)

        self.Vh_out = nn.Parameter(
            torch.randn(self.rank_out, self.all_head_size) * 0.02, requires_grad=False
        )
        self.S_out_sqrt = nn.Parameter(torch.ones(self.rank_out), requires_grad=False)
        self.U_out = nn.Parameter(
            torch.randn(self.hidden_size, self.rank_out) * 0.02, requires_grad=False
        )
        self.bias_out = nn.Parameter(torch.zeros(self.hidden_size), requires_grad=False)

        # Trainable LoRA residual adapters
        if self.lora_rank > 0:
            self.lora_A_qkv = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.lora_B_qkv = nn.Parameter(torch.empty(qkv_out_dim, self.lora_rank))
            self.lora_A_out = nn.Parameter(torch.empty(self.lora_rank, self.all_head_size))
            self.lora_B_out = nn.Parameter(torch.empty(self.hidden_size, self.lora_rank))
            self.reset_lora_parameters()
        else:
            self.register_parameter("lora_A_qkv", None)
            self.register_parameter("lora_B_qkv", None)
            self.register_parameter("lora_A_out", None)
            self.register_parameter("lora_B_out", None)

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", 0.1))
            if cfg
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

        src = pretrained_attn or (
            original_layer
            if hasattr(original_layer, "q_lin")
            or hasattr(original_layer, "query")
            or hasattr(original_layer, "c_attn")
            else None
        )
        if src is not None:
            self.copy_weights(src)

    @property
    def qkv_weight(self) -> torch.Tensor:
        """Reconstruct the effective base QKV weight from SVD factors."""
        return (self.U_qkv * self.S_qkv_sqrt.view(1, -1)) @ (
            self.Vh_qkv * self.S_qkv_sqrt.view(-1, 1)
        )

    @property
    def out_weight(self) -> torch.Tensor:
        """Reconstruct the effective base Out weight from SVD factors."""
        return (self.U_out * self.S_out_sqrt.view(1, -1)) @ (
            self.Vh_out * self.S_out_sqrt.view(-1, 1)
        )

    def reset_lora_parameters(self) -> None:
        if self.lora_rank > 0:
            nn.init.kaiming_uniform_(self.lora_A_qkv, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_qkv)
            nn.init.kaiming_uniform_(self.lora_A_out, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_out)
            self.lora_A_qkv.requires_grad = True
            self.lora_B_qkv.requires_grad = True
            self.lora_A_out.requires_grad = True
            self.lora_B_out.requires_grad = True

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. HuggingFace / standard 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 as e:
                logger.debug(
                    f"UniversalAdapter failed for {proj}: {e}"
                )  # continue to fallback paths

        # 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
        ):
            weight = src.c_attn.weight.data
            bias = getattr(src.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                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:
            w_o = src.c_proj.weight.data
            if (
                w_o.dim() == 2
                and w_o.shape[0] != self.hidden_size
                and w_o.shape[1] == self.hidden_size
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(src.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        # Check output projection attrs FIRST — DistilBERT has out_lin and is NOT nested
        if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
            self.is_nested_encoder = False
        elif (
            (
                hasattr(src, "self")
                and hasattr(src.self, "query")
                and hasattr(src, "output")
                and hasattr(src.output, "dense")
            )
            or (
                cls_name in ("BertSelfAttention", "RobertaSelfAttention")
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        # Check required weights
        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(
                f"CompressedAttentionLayer: missing weights for projections {missing}; proceeding with random init"
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"CompressedAttentionLayer: extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f}); continuing"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"CompressedAttentionLayer: extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f}); continuing"
                )

        # Joint SVD on Q, K, V
        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():
            w_q = _orient(extracted["q"][0].float(), self.hidden_size)
            w_k = _orient(extracted["k"][0].float(), self.hidden_size)
            w_v = _orient(extracted["v"][0].float(), self.hidden_size)
            w_qkv = torch.cat([w_q, w_k, w_v], dim=0)

            U, S, Vh = torch.linalg.svd(w_qkv, full_matrices=False)
            r_qkv = min(self.rank_qkv, U.size(1), Vh.size(0))
            sqrt_s = torch.sqrt(torch.clamp(S[:r_qkv], min=1e-12))

            qkv_out_dim = self.all_head_size + 2 * self.kv_head_size
            self.U_qkv.data[: min(qkv_out_dim, U.size(0)), :r_qkv].copy_(
                U[: min(qkv_out_dim, U.size(0)), :r_qkv]
            )
            self.S_qkv_sqrt.data[:r_qkv].copy_(sqrt_s)
            self.Vh_qkv.data[:r_qkv, : min(self.hidden_size, Vh.size(1))].copy_(
                Vh[:r_qkv, : min(self.hidden_size, Vh.size(1))]
            )

            b_q = (
                extracted["q"][1]
                if extracted["q"][1] is not None
                else torch.zeros(w_q.size(0), device=w_q.device)
            )
            b_k = (
                extracted["k"][1]
                if extracted["k"][1] is not None
                else torch.zeros(w_k.size(0), device=w_k.device)
            )
            b_v = (
                extracted["v"][1]
                if extracted["v"][1] is not None
                else torch.zeros(w_v.size(0), device=w_v.device)
            )
            b_q_val = b_q.data if isinstance(b_q, nn.Parameter) else b_q
            b_k_val = b_k.data if isinstance(b_k, nn.Parameter) else b_k
            b_v_val = b_v.data if isinstance(b_v, nn.Parameter) else b_v
            b_qkv = torch.cat([b_q_val, b_k_val, b_v_val], dim=0)
            self.bias_qkv.data[: min(qkv_out_dim, b_qkv.size(0))].copy_(
                b_qkv[: min(qkv_out_dim, b_qkv.size(0))]
            )

            # Output projection SVD
            if extracted["o"][0] is not None:
                w_o = _orient(extracted["o"][0].float(), self.all_head_size)
                U_o, S_o, Vh_o = torch.linalg.svd(w_o, full_matrices=False)
                r_out = min(self.rank_out, U_o.size(1), Vh_o.size(0))
                sqrt_so = torch.sqrt(torch.clamp(S_o[:r_out], min=1e-12))

                self.U_out.data[: min(self.hidden_size, U_o.size(0)), :r_out].copy_(
                    U_o[: min(self.hidden_size, U_o.size(0)), :r_out]
                )
                self.S_out_sqrt.data[:r_out].copy_(sqrt_so)
                self.Vh_out.data[:r_out, : min(self.all_head_size, Vh_o.size(1))].copy_(
                    Vh_o[:r_out, : min(self.all_head_size, Vh_o.size(1))]
                )

                if extracted["o"][1] is not None:
                    b_o = extracted["o"][1]
                    b_o_data = b_o.data if isinstance(b_o, nn.Parameter) else b_o
                    self.bias_out.data[: min(self.hidden_size, b_o_data.size(0))].copy_(
                        b_o_data[: min(self.hidden_size, b_o_data.size(0))]
                    )
            elif self.is_nested_encoder:
                self.U_out.data.zero_()
                min_d = min(self.hidden_size, self.rank_out)
                self.U_out.data[:min_d, :min_d].fill_diagonal_(1.0)
                self.S_out_sqrt.data.fill_(1.0)
                self.Vh_out.data.zero_()
                self.Vh_out.data[:min_d, :min_d].fill_diagonal_(1.0)
                self.bias_out.data.zero_()

            # Ensure base weights are frozen
            self.Vh_qkv.requires_grad = False
            self.S_qkv_sqrt.requires_grad = False
            self.U_qkv.requires_grad = False
            self.bias_qkv.requires_grad = False
            self.Vh_out.requires_grad = False
            self.S_out_sqrt.requires_grad = False
            self.U_out.requires_grad = False
            self.bias_out.requires_grad = False

            if self.lora_rank > 0:
                self.reset_lora_parameters()

        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]:
        if hidden_states.dim() == 2:
            hidden_states = hidden_states.unsqueeze(0)
        B, N, D = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        # FIX 1: Pre-transpose weight matrices to avoid runtime transposes
        # Vh_qkv is (Rank, D), U_qkv is (Out_Dim, Rank)
        # S_qkv_sqrt is (Rank)

        # Optimize Matmul path
        v_proj = torch.matmul(hidden_states, self.Vh_qkv.t()) * self.S_qkv_sqrt
        qkv = torch.matmul(v_proj, (self.U_qkv * self.S_qkv_sqrt).t()) + self.bias_qkv

        # LoRA residual adapter
        if self.lora_rank > 0 and self.lora_A_qkv is not None and self.lora_B_qkv is not None:
            lora_qkv = (
                F.linear(F.linear(hidden_states, self.lora_A_qkv), self.lora_B_qkv)
                * self.lora_scaling
            )
            qkv = qkv + lora_qkv

        # Split into Q, K, V
        q, k, v = torch.split(
            qkv, [self.all_head_size, self.kv_head_size, self.kv_head_size], dim=-1
        )

        # FIX 2: Remove .contiguous() memory bombs
        q = q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # Apply RoPE
        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"))

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

        scale = 1.0 / math.sqrt(self.head_dim)

        # Use torch.baddbmm or scaled_dot_product_attention for huge memory savings
        # Fallback to standard matmul if not available or unsupported shapes
        try:
            # PyTorch 2.0+ Flash Attention / Memory Efficient Attention
            # This handles Causal masking automatically and saves massive VRAM
            is_causal = self.is_decoder and (attention_mask is None or attention_mask.dim() != 4)
            attn_out = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=attention_mask,
                dropout_p=self.dropout.p if self.training else 0.0,
                is_causal=is_causal,
            )
            probs = None  # SDPA doesn't return probabilities to save memory
        except Exception:
            # Legacy Path (if SDPA fails due to old PyTorch or weird masks)
            scores = torch.matmul(q, k.transpose(-2, -1)) * scale

            mask = attention_mask if attention_mask is not None else kwargs.get("mask")
            if mask is not None:
                if mask.dim() == 4:
                    if mask.dtype == torch.bool:
                        scores = scores.masked_fill(~mask, -10000.0)
                    else:
                        scores = scores + mask.to(scores.dtype)
                else:
                    valid_m = get_token_valid_mask(mask, B, N, device)
                    if valid_m is not None:
                        scores = scores.masked_fill((~valid_m).unsqueeze(1).unsqueeze(2), -10000.0)

            # Explicit Causal Masking
            if self.is_decoder and (mask is None or mask.dim() != 4):
                causal_tril = (
                    torch.tril(torch.ones((N, N), device=device, dtype=torch.bool))
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores.masked_fill(~causal_tril, -10000.0)

            # FIX 3: Safe Softmax without extra type casts that create new tensors
            probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(dtype)
            probs = self.dropout(probs)
            attn_out = torch.matmul(probs, v)

        # Transpose back and view (reshape handles non-contiguous memory safely without deep copies)
        attn_out = attn_out.transpose(1, 2).reshape(B, N, self.all_head_size)

        if self.is_nested_encoder:
            out = attn_out
        else:
            # SVD Out forward
            out_v = torch.matmul(attn_out, self.Vh_out.t()) * self.S_out_sqrt
            out = torch.matmul(out_v, (self.U_out * self.S_out_sqrt).t()) + self.bias_out
            if self.lora_rank > 0 and self.lora_A_out is not None and self.lora_B_out is not None:
                lora_out = (
                    F.linear(F.linear(attn_out, self.lora_A_out), self.lora_B_out)
                    * self.lora_scaling
                )
                out = out + lora_out

        return out, (probs if output_attentions and probs is not None else None)

Every 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

Subclasses

• CompressedAttentionLayer2

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 out_weight : torch.Tensor

Expand source code

@property
def out_weight(self) -> torch.Tensor:
    """Reconstruct the effective base Out weight from SVD factors."""
    return (self.U_out * self.S_out_sqrt.view(1, -1)) @ (
        self.Vh_out * self.S_out_sqrt.view(-1, 1)
    )

Reconstruct the effective base Out weight from SVD factors.

prop qkv_weight : torch.Tensor

Expand source code

@property
def qkv_weight(self) -> torch.Tensor:
    """Reconstruct the effective base QKV weight from SVD factors."""
    return (self.U_qkv * self.S_qkv_sqrt.view(1, -1)) @ (
        self.Vh_qkv * self.S_qkv_sqrt.view(-1, 1)
    )

Reconstruct the effective base QKV weight from SVD factors.

Methods

def reset_lora_parameters(self) ‑> None

Expand source code

def reset_lora_parameters(self) -> None:
    if self.lora_rank > 0:
        nn.init.kaiming_uniform_(self.lora_A_qkv, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B_qkv)
        nn.init.kaiming_uniform_(self.lora_A_out, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B_out)
        self.lora_A_qkv.requires_grad = True
        self.lora_B_qkv.requires_grad = True
        self.lora_A_out.requires_grad = True
        self.lora_B_out.requires_grad = True

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class CompressedAttentionLayer2 (config: Any | None = None,
rank: int | None = None,
rank_qkv: int = 256,
rank_out: int = 256,
lora_rank: int = 16,
original_layer: Any = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("compressed_attention_layer_energy")
class CompressedAttentionLayer2(CompressedAttentionLayer):
    pass

Every custom attention module must inherit from this.

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

Ancestors

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

Inherited members

• CompressedAttentionLayer:
  • copy_weights
  • extra_repr
  • forward
  • is_custom_attention
  • out_weight
  • qkv_weight
  • supports_bidirectional
  • supports_causal
  • supports_gqa
  • supports_rope

class CompressedDistilBertAttention (model_config: Any | None = None,
rank: int | None = None,
rank_qkv: int = 256,
rank_out: int = 256,
dropout: float = 0.1,
lora_rank: int = 16,
pretrained_attn: nn.Module | None = None,
config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "compressed_mha",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "compressed_mha_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class CompressedDistilBertAttention(BaseAttention):
    supports_causal: bool = True
    supports_bidirectional: bool = True
    supports_gqa: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        model_config: Any | None = None,
        rank: int | None = None,
        rank_qkv: int = 256,
        rank_out: int = 256,
        dropout: float = 0.1,
        lora_rank: int = 16,
        pretrained_attn: nn.Module | None = None,
        config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config or model_config
        super().__init__(cfg, rank)
        self.adapter = ModelAdapter(cfg)
        self.hidden_size: int = self.adapter.hidden_size or kwargs.get(
            "hidden_size", kwargs.get("d_model", 768)
        )
        self.num_heads: int = self.adapter.num_heads or kwargs.get(
            "num_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads: int = self.adapter.num_key_value_heads or kwargs.get(
            "num_kv_heads", kwargs.get("num_key_value_heads", self.num_heads)
        )
        self.head_dim: int = self.adapter.head_dim or (self.hidden_size // max(1, self.num_heads))
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))
        _mt = getattr(cfg, "model_type", "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        self.is_decoder: bool = bool(
            getattr(cfg, "is_decoder", False) or any(_mt.startswith(p) for p in _decoder_prefixes)
        )
        self.is_nested_encoder: bool = False

        rank_val = rank or kwargs.get("rank_qkv", rank_qkv)
        self.rank_qkv = min(rank_val, self.all_head_size + 2 * self.kv_head_size, self.hidden_size)
        out_r = rank or kwargs.get("rank_out", rank_out)
        self.rank_out = min(out_r, self.all_head_size, self.hidden_size)
        self.lora_rank = getattr(cfg, "lora_rank", lora_rank)
        self.lora_scaling = (32.0 / self.lora_rank) if self.lora_rank > 0 else 1.0

        qkv_out_dim = self.all_head_size + 2 * self.kv_head_size

        # Frozen base SVD parameters
        self.U_qkv = nn.Parameter(
            torch.randn(qkv_out_dim, self.rank_qkv) * 0.02, requires_grad=False
        )
        self.S_qkv = nn.Parameter(torch.ones(self.rank_qkv), requires_grad=False)
        self.Vh_qkv = nn.Parameter(
            torch.randn(self.rank_qkv, self.hidden_size) * 0.02, requires_grad=False
        )

        self.U_out = nn.Parameter(
            torch.randn(self.hidden_size, self.rank_out) * 0.02, requires_grad=False
        )
        self.S_out = nn.Parameter(torch.ones(self.rank_out), requires_grad=False)
        self.Vh_out = nn.Parameter(
            torch.randn(self.rank_out, self.all_head_size) * 0.02, requires_grad=False
        )

        self.bias_q = nn.Parameter(torch.zeros(self.all_head_size), requires_grad=False)
        self.bias_k = nn.Parameter(torch.zeros(self.kv_head_size), requires_grad=False)
        self.bias_v = nn.Parameter(torch.zeros(self.kv_head_size), requires_grad=False)
        self.bias_out = nn.Parameter(torch.zeros(self.hidden_size), requires_grad=False)

        # Trainable LoRA residual adapters
        if self.lora_rank > 0:
            self.lora_A_qkv = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.lora_B_qkv = nn.Parameter(torch.empty(qkv_out_dim, self.lora_rank))
            self.lora_A_out = nn.Parameter(torch.empty(self.lora_rank, self.all_head_size))
            self.lora_B_out = nn.Parameter(torch.empty(self.hidden_size, self.lora_rank))
            self.reset_lora_parameters()
        else:
            self.register_parameter("lora_A_qkv", None)
            self.register_parameter("lora_B_qkv", None)
            self.register_parameter("lora_A_out", None)
            self.register_parameter("lora_B_out", None)

        self.dropout_p = getattr(cfg, "dropout", dropout) if cfg else dropout
        self.dropout = nn.Dropout(self.dropout_p)
        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

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

    @property
    def qkv_weight(self) -> torch.Tensor:
        """Reconstruct the effective base QKV weight from SVD factors."""
        return (self.U_qkv * self.S_qkv.view(1, -1)) @ self.Vh_qkv

    @property
    def out_weight(self) -> torch.Tensor:
        """Reconstruct the effective base Out weight from SVD factors."""
        return (self.U_out * self.S_out.view(1, -1)) @ self.Vh_out

    def reset_lora_parameters(self) -> None:
        if self.lora_rank > 0:
            nn.init.kaiming_uniform_(self.lora_A_qkv, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_qkv)
            nn.init.kaiming_uniform_(self.lora_A_out, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B_out)
            self.lora_A_qkv.requires_grad = True
            self.lora_B_qkv.requires_grad = True
            self.lora_A_out.requires_grad = True
            self.lora_B_out.requires_grad = True

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. HuggingFace / standard 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 as e:
                logger.debug(
                    f"UniversalAdapter failed for {proj}: {e}"
                )  # continue to fallback paths

        # 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
        ):
            weight = src.c_attn.weight.data
            bias = getattr(src.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                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:
            w_o = src.c_proj.weight.data
            p_cls = getattr(src.c_proj, "__class__", type(src.c_proj)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(src.c_proj, "nf")
                or (
                    w_o.dim() == 2
                    and (
                        w_o.shape[0] < w_o.shape[1]
                        or (w_o.shape[0] == self.hidden_size and w_o.shape[1] == self.hidden_size)
                    )
                )
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(src.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        # Check output projection attrs FIRST — DistilBERT has out_lin and is NOT nested
        if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
            self.is_nested_encoder = False
        elif (
            (
                hasattr(src, "self")
                and hasattr(src.self, "query")
                and hasattr(src, "output")
                and hasattr(src.output, "dense")
            )
            or (
                cls_name in ("BertSelfAttention", "RobertaSelfAttention")
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        # Check required weights
        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(
                f"CompressedDistilBertAttention: missing weights for projections {missing}; proceeding with random init"
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"CompressedDistilBertAttention: extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f}); continuing"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"CompressedDistilBertAttention: extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f}); continuing"
                )

        # Joint SVD on Q, K, V
        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():
            w_q = _orient(extracted["q"][0].float(), self.hidden_size)
            w_k = _orient(extracted["k"][0].float(), self.hidden_size)
            w_v = _orient(extracted["v"][0].float(), self.hidden_size)
            w_qkv = torch.cat([w_q, w_k, w_v], dim=0)

            U, S, Vh = torch.linalg.svd(w_qkv, full_matrices=False)
            r_qkv = min(self.rank_qkv, U.size(1), Vh.size(0))

            qkv_out_dim = self.all_head_size + 2 * self.kv_head_size
            self.U_qkv.data[: min(qkv_out_dim, U.size(0)), :r_qkv].copy_(
                U[: min(qkv_out_dim, U.size(0)), :r_qkv]
            )
            self.S_qkv.data[:r_qkv].copy_(S[:r_qkv])
            self.Vh_qkv.data[:r_qkv, : min(self.hidden_size, Vh.size(1))].copy_(
                Vh[:r_qkv, : min(self.hidden_size, Vh.size(1))]
            )

            if extracted["q"][1] is not None:
                b_q = extracted["q"][1]
                b_q_data = b_q.data if isinstance(b_q, nn.Parameter) else b_q
                self.bias_q.data[: min(self.all_head_size, len(b_q_data))].copy_(
                    b_q_data[: min(self.all_head_size, len(b_q_data))]
                )
            if extracted["k"][1] is not None:
                b_k = extracted["k"][1]
                b_k_data = b_k.data if isinstance(b_k, nn.Parameter) else b_k
                self.bias_k.data[: min(self.kv_head_size, len(b_k_data))].copy_(
                    b_k_data[: min(self.kv_head_size, len(b_k_data))]
                )
            if extracted["v"][1] is not None:
                b_v = extracted["v"][1]
                b_v_data = b_v.data if isinstance(b_v, nn.Parameter) else b_v
                self.bias_v.data[: min(self.kv_head_size, len(b_v_data))].copy_(
                    b_v_data[: min(self.kv_head_size, len(b_v_data))]
                )

            # Output projection SVD
            if extracted["o"][0] is not None:
                w_o = _orient(extracted["o"][0].float(), self.all_head_size)
                U_o, S_o, Vh_o = torch.linalg.svd(w_o, full_matrices=False)
                r_out = min(self.rank_out, U_o.size(1), Vh_o.size(0))

                self.U_out.data[: min(self.hidden_size, U_o.size(0)), :r_out].copy_(
                    U_o[: min(self.hidden_size, U_o.size(0)), :r_out]
                )
                self.S_out.data[:r_out].copy_(S_o[:r_out])
                self.Vh_out.data[:r_out, : min(self.all_head_size, Vh_o.size(1))].copy_(
                    Vh_o[:r_out, : min(self.all_head_size, Vh_o.size(1))]
                )

                if extracted["o"][1] is not None:
                    b_o = extracted["o"][1]
                    b_o_data = b_o.data if isinstance(b_o, nn.Parameter) else b_o
                    self.bias_out.data[: min(self.hidden_size, len(b_o_data))].copy_(
                        b_o_data[: min(self.hidden_size, len(b_o_data))]
                    )
            elif self.is_nested_encoder:
                self.U_out.data.zero_()
                min_d = min(self.hidden_size, self.rank_out)
                self.U_out.data[:min_d, :min_d].fill_diagonal_(1.0)
                self.S_out.data.fill_(1.0)
                self.Vh_out.data.zero_()
                self.Vh_out.data[:min_d, :min_d].fill_diagonal_(1.0)
                self.bias_out.data.zero_()

            # Ensure base weights are frozen
            self.U_qkv.requires_grad = False
            self.S_qkv.requires_grad = False
            self.Vh_qkv.requires_grad = False
            self.U_out.requires_grad = False
            self.S_out.requires_grad = False
            self.Vh_out.requires_grad = False
            self.bias_q.requires_grad = False
            self.bias_k.requires_grad = False
            self.bias_v.requires_grad = False
            self.bias_out.requires_grad = False

            if self.lora_rank > 0:
                self.reset_lora_parameters()

        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]:
        if (
            attention_mask is not None
            and attention_mask.dtype == torch.bool
            and attention_mask.shape[-1] == self.num_heads
        ):
            head_mask, attention_mask = attention_mask, None
        if hidden_states.dim() == 2:
            hidden_states = hidden_states.unsqueeze(0)
        B, N, _ = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        proj = torch.matmul(hidden_states, self.Vh_qkv.t().contiguous()) * self.S_qkv.unsqueeze(
            0
        ).unsqueeze(0)
        qkv_comp = torch.matmul(proj, self.U_qkv.t().contiguous())

        # LoRA residual adapter
        if self.lora_rank > 0 and self.lora_A_qkv is not None and self.lora_B_qkv is not None:
            lora_qkv = (
                F.linear(F.linear(hidden_states, self.lora_A_qkv), self.lora_B_qkv)
                * self.lora_scaling
            )
            qkv_comp = qkv_comp + lora_qkv

        q_raw, k_raw, v_raw = torch.split(
            qkv_comp, [self.all_head_size, self.kv_head_size, self.kv_head_size], dim=-1
        )
        q = q_raw + self.bias_q
        k = k_raw + self.bias_k
        v = v_raw + self.bias_v

        q = q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k = k.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(B, N, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()

        # Apply RoPE
        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"))

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

        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1).contiguous()) * scale

        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, device)
                if valid_m is not None:
                    scores = scores.masked_fill((~valid_m).unsqueeze(1).unsqueeze(2), -10000.0)

        # Explicit Causal Masking for decoders
        mt = str(getattr(self.config, "model_type", "") or "").lower()
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_decoder = (
            self.is_decoder
            or getattr(self.config, "is_decoder", False)
            or any(mt.startswith(p) for p in decoder_prefixes)
        )
        if is_decoder:
            causal_tril = (
                torch.tril(torch.ones((N, N), device=device, dtype=torch.bool))
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores.masked_fill(~causal_tril, -10000.0)

        attn_probs = torch.nan_to_num(F.softmax(scores.float(), dim=-1), nan=0.0).to(dtype)
        attn_probs = self.dropout(attn_probs)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn_probs = attn_probs * head_mask

        context = torch.matmul(attn_probs, v)
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            out = context
        else:
            out_proj = torch.matmul(context, self.Vh_out.t().contiguous()) * self.S_out.unsqueeze(
                0
            ).unsqueeze(0)
            out = torch.matmul(out_proj, self.U_out.t().contiguous()) + self.bias_out
            if self.lora_rank > 0 and self.lora_A_out is not None and self.lora_B_out is not None:
                lora_out = (
                    F.linear(F.linear(context, self.lora_A_out), self.lora_B_out)
                    * self.lora_scaling
                )
                out = out + lora_out

        return out, (attn_probs if output_attentions else None)

Every 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_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 out_weight : torch.Tensor

Expand source code

@property
def out_weight(self) -> torch.Tensor:
    """Reconstruct the effective base Out weight from SVD factors."""
    return (self.U_out * self.S_out.view(1, -1)) @ self.Vh_out

Reconstruct the effective base Out weight from SVD factors.

prop qkv_weight : torch.Tensor

Expand source code

@property
def qkv_weight(self) -> torch.Tensor:
    """Reconstruct the effective base QKV weight from SVD factors."""
    return (self.U_qkv * self.S_qkv.view(1, -1)) @ self.Vh_qkv

Reconstruct the effective base QKV weight from SVD factors.

Methods

def reset_lora_parameters(self) ‑> None

Expand source code

def reset_lora_parameters(self) -> None:
    if self.lora_rank > 0:
        nn.init.kaiming_uniform_(self.lora_A_qkv, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B_qkv)
        nn.init.kaiming_uniform_(self.lora_A_out, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B_out)
        self.lora_A_qkv.requires_grad = True
        self.lora_B_qkv.requires_grad = True
        self.lora_A_out.requires_grad = True
        self.lora_B_out.requires_grad = True

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class ContextFusionGate (hidden_size, num_heads, head_dim, init_bias=1.8)

Expand source code

class ContextFusionGate(nn.Module):
    def __init__(self, hidden_size, num_heads, head_dim, init_bias=1.8):
        super().__init__()
        self.hidden_proj = (
            nn.Linear(hidden_size, num_heads * head_dim, bias=False)
            if hidden_size != num_heads * head_dim
            else None
        )
        self.proj = nn.Linear(3 * head_dim, 1, bias=True)
        nn.init.constant_(self.proj.bias, init_bias)

    def forward(self, hidden_states, local_out, global_out):
        B, h, N, d = local_out.shape
        h_exp = self.hidden_proj(hidden_states) if self.hidden_proj is not None else hidden_states
        h_per_head = h_exp.view(B, N, h, d).permute(0, 2, 1, 3)
        return torch.sigmoid(self.proj(torch.cat([h_per_head, local_out, global_out], dim=-1)))

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, hidden_states, local_out, global_out) ‑> Callable[..., typing.Any]

Expand source code

def forward(self, hidden_states, local_out, global_out):
    B, h, N, d = local_out.shape
    h_exp = self.hidden_proj(hidden_states) if self.hidden_proj is not None else hidden_states
    h_per_head = h_exp.view(B, N, h, d).permute(0, 2, 1, 3)
    return torch.sigmoid(self.proj(torch.cat([h_per_head, local_out, global_out], dim=-1)))

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class CosformerAttention (config: Any | None = None,
rank: int | None = None,
max_seq_len: int = 512,
act_fun: str = 'relu',
eps: float = 1e-06,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "cosformer", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "cosformer_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class CosformerAttention(BaseAttention):
    r"""Cosformer Attention with Linear Complexity and Exact Multi-Architecture Support."""

    is_custom_attention: bool = True
    triton_kernel = "linear_attn"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        max_seq_len: int = 512,
        act_fun: str = "relu",
        eps: float = 1e-6,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        def _cfg_get(k: str, default: Any = None) -> Any:
            for obj in (config, comp_config):
                if obj is not None and hasattr(obj, k):
                    val = getattr(obj, k)
                    if val is not None:
                        return val
            return kwargs.get(k, default)

        hidden_size = kwargs.get("hidden_size")
        num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")

        if hidden_size is None:
            hidden_size = _cfg_get("hidden_size", _cfg_get("dim", None))
            if hidden_size is None and source_attn is not None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                hidden_size = (
                    getattr(out_layer, "out_features", 768) if out_layer is not None else 768
                )
            elif hidden_size is None:
                hidden_size = 768

        if num_heads is None:
            num_heads = _cfg_get("num_attention_heads", _cfg_get("num_heads", None))
            if num_heads is None and source_attn is not None:
                num_heads = getattr(
                    source_attn,
                    "n_heads",
                    getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(source_attn, "num_heads", 12),
                    ),
                )
            elif num_heads is None:
                num_heads = 12

        super().__init__(config, rank)
        self.dim: int = int(hidden_size)
        self.hidden_size: int = self.dim
        self.num_heads: int = int(num_heads)

        resolved_head_dim = _cfg_get("head_dim", None)
        if resolved_head_dim is None and source_attn is not None:
            resolved_head_dim = getattr(source_attn, "head_dim", None)
        if resolved_head_dim is None:
            resolved_head_dim = self.dim // max(1, self.num_heads)
        self.head_dim: int = int(resolved_head_dim)

        # Fix 11: even head_dim required for RoPE
        if self.head_dim % 2 != 0:
            raise ValueError(
                f"CosformerAttention with RoPE requires even head_dim, got head_dim={self.head_dim}"
            )

        kv_heads = _cfg_get("num_key_value_heads", None)
        if kv_heads is None and source_attn is not None:
            kv_heads = getattr(
                source_attn,
                "num_key_value_heads",
                getattr(
                    source_attn,
                    "num_kv_heads",
                    getattr(source_attn, "n_kv_heads", None),
                ),
            )
            if (
                kv_heads is None
                and hasattr(source_attn, "k_proj")
                and hasattr(source_attn.k_proj, "out_features")
                and self.head_dim > 0
            ):
                kv_heads = source_attn.k_proj.out_features // self.head_dim
        self.num_kv_heads: int = int(kv_heads if kv_heads is not None else self.num_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
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))

        # Fix 9: documented max_seq_len default
        self.max_seq_len: int = int(_cfg_get("max_position_embeddings", max_seq_len))
        self.act_fun: str = act_fun
        self.eps: float = float(eps)

        # Static layer index
        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx: int = int(resolved_layer_id)

        # LoRA adaptation configuration
        self.lora_rank: int = int(_cfg_get("lora_rank", kwargs.get("lora_rank", 16)))
        self.lora_alpha: float = float(_cfg_get("lora_alpha", kwargs.get("lora_alpha", 32.0)))
        self.freeze_base: bool = self.lora_rank > 0

        # Decoder & nested encoder detection
        model_type = getattr(config, "model_type", "").lower() if config is not None else ""
        is_dec = _cfg_get("is_decoder", _cfg_get("is_causal", kwargs.get("is_causal", False)))
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "phi",
            "gemma",
            "opt",
            "bloom",
            "falcon",
        )
        if (
            any(model_type.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)
        self.is_nested_encoder: bool = bool(kwargs.get("is_nested_encoder", False))

        # Multi-head projections equipped with LoRA adapters
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        self.rotary_emb = None
        self.comp_config: Any | None = comp_config
        self.triton_kwargs = {
            "phi": "elu+1",
            "eps": self.eps,
            "cos_weight": None,
            "sin_weight": None,
        }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        """Transfer pretrained projection weights with strict Frobenius norm verification."""
        if src is None:
            return

        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")
        )
        if cls_name in ["BertSelfAttention", "RobertaSelfAttention"] or (
            hasattr(src, "query") and not has_out
        ):
            self.is_nested_encoder = True
        if hasattr(src, "self") and hasattr(src, "output"):
            self.is_nested_encoder = True

        if hasattr(src, "rotary_emb"):
            self.rotary_emb = src.rotary_emb

        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 + c_proj
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c_attn_mod = src.c_attn
            if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    logger.debug(f"c_attn_mod.merge failed: {e}")
            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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (c_proj_mod.__class__.__name__ == "Conv1D")
                        or (base_o.__class__.__name__ == "Conv1D")
                        or hasattr(c_proj_mod, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        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. RoBERTa / BERT nested self-attention
        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)
                base_mod = getattr(mod, "base_layer", mod)
                if base_mod is not None and hasattr(base_mod, "weight"):
                    extracted[proj_key] = (
                        base_mod.weight.data.clone(),
                        (
                            base_mod.bias.data.clone()
                            if getattr(base_mod, "bias", None) is not None
                            else None
                        ),
                    )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                out_mod = src.output.dense
                base_out = getattr(out_mod, "base_layer", out_mod)
                extracted["o"] = (
                    base_out.weight.data.clone(),
                    (
                        base_out.bias.data.clone()
                        if getattr(base_out, "bias", None) is not None
                        else None
                    ),
                )

        # 3. Standard linear projections (DistilBERT, Qwen, LLaMA)
        else:
            proj_candidates = {
                "q": ["q_lin", "q_proj", "query", "W_q"],
                "k": ["k_lin", "k_proj", "key", "W_k"],
                "v": ["v_lin", "v_proj", "value", "W_v"],
                "o": ["out_lin", "out_proj", "dense", "o_proj", "c_proj", "wo"],
            }
            for proj_key, names in proj_candidates.items():
                for name in names:
                    if hasattr(src, name):
                        mod = getattr(src, name)
                        base_mod = getattr(mod, "base_layer", mod)
                        if hasattr(base_mod, "weight"):
                            w_data = base_mod.weight.data.clone()
                            if (
                                (base_mod.__class__.__name__ == "Conv1D")
                                or hasattr(base_mod, "nf")
                                or (
                                    w_data.ndim == 2
                                    and proj_key == "o"
                                    and w_data.shape[0] < w_data.shape[1]
                                )
                            ):
                                w_data = w_data.t().contiguous()
                            b_data = getattr(base_mod, "bias", None)
                            extracted[proj_key] = (
                                w_data,
                                b_data.data.clone() if b_data is not None else None,
                            )
                            break

        # Fallback via UniversalAdapter
        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(f"UniversalAdapter fallback: {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"CosformerAttention copy_weights: missing mandatory projections "
                f"{missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

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

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    # ------------------------------------------------------------------
    # Positional weighting
    # ------------------------------------------------------------------
    def _get_cos_sin_weights(
        self, seq_len: int, device: torch.device, dtype: torch.dtype
    ) -> tuple[torch.Tensor, torch.Tensor]:
        r"""Generate positional cosine and sine factors: [1, 1, seq_len, 1].

        Fix 6: Uses seq_len directly (simpler and more numerically stable).
        """
        pos = torch.arange(seq_len, device=device, dtype=torch.float32)
        m = max(1, int(seq_len))
        angles = (math.pi / (2.0 * m)) * pos
        cos_weight = torch.cos(angles).view(1, 1, seq_len, 1).to(dtype=dtype)
        sin_weight = torch.sin(angles).view(1, 1, seq_len, 1).to(dtype=dtype)
        return cos_weight, sin_weight

    # ------------------------------------------------------------------
    # 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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        # ------------------------------------------------------------------
        # 0. Input validation
        # ------------------------------------------------------------------
        assert hidden_states.dim() == 3, (
            f"CosformerAttention expects 3D hidden_states [B, N, D], "
            f"got shape {tuple(hidden_states.shape)}"
        )

        B, N, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        input_dtype = hidden_states.dtype

        is_cross_attention = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross_attention else hidden_states
        kv_len = kv_input.shape[1]

        # ------------------------------------------------------------------
        # 1. Q/K/V projections
        # ------------------------------------------------------------------
        Q = self.q_proj(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # ------------------------------------------------------------------
        # 2. RoPE (only for self-attention, not cross-attention)
        # ------------------------------------------------------------------
        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        if pos_emb is not None and not is_cross_attention:
            if isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
                cos, sin = pos_emb[0], pos_emb[1]
                Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids)
            else:
                logger.warning(
                    "CosformerAttention: unexpected position_embeddings format (%s), skipping RoPE.",
                    type(pos_emb),
                )
        elif self.rotary_emb is not None and not is_cross_attention:
            rotary_mod = self.rotary_emb
            pos_ids = position_ids if position_ids is not None else kwargs.get("position_ids")
            if pos_ids is None:
                pos_ids = torch.arange(N, device=device).unsqueeze(0).expand(B, -1)
            cos, sin = None, None
            try:
                cos, sin = rotary_mod(Q, seq_len=kv_len)
            except Exception:
                try:
                    cos, sin = rotary_mod(Q, pos_ids)
                except Exception:
                    try:
                        cos, sin = rotary_mod(Q)
                    except Exception:
                        pass
            if cos is not None and sin is not None:
                Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids=pos_ids)

        # ------------------------------------------------------------------
        # 3. GQA head expansion
        # ------------------------------------------------------------------
        if self.num_kv_heads != self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        # ------------------------------------------------------------------
        # 4. Kernel feature mapping (ELU+1) — in FP32 for stability
        # ------------------------------------------------------------------
        if self.act_fun == "softplus":
            Q_k = F.softplus(Q.float())
            K_k = F.softplus(K.float())
        else:
            # default: elu + 1 → non-negative kernel feature map
            Q_k = F.elu(Q.float()) + 1.0
            K_k = F.elu(K.float()) + 1.0
        V_f = V.float()

        # ------------------------------------------------------------------
        # 5. Token valid mask on both K and V (padding tokens → 0)
        # ------------------------------------------------------------------
        eff_mask = encoder_attention_mask if is_cross_attention else attention_mask
        valid_mask = None
        if eff_mask is not None:
            valid_mask = get_token_valid_mask(eff_mask, B, kv_len, device)  # [B, kv_len] bool
            mask_expanded = valid_mask.unsqueeze(1).unsqueeze(-1).float()  # [B, 1, kv_len, 1]
            K_k = K_k * mask_expanded
            V_f = V_f * mask_expanded

        # ------------------------------------------------------------------
        # 6. Cos/Sin positional reweighting (cosformer's signature)
        # ------------------------------------------------------------------
        cos_w, sin_w = self._get_cos_sin_weights(N, device, dtype=torch.float32)
        # cos_w, sin_w: [1, 1, N, 1]
        Q_cos = Q_k * cos_w
        Q_sin = Q_k * sin_w
        K_cos = K_k * cos_w
        K_sin = K_k * sin_w

        # ------------------------------------------------------------------
        # 7. Linear attention — try Triton, fallback to stable Python path
        # ------------------------------------------------------------------
        is_causal_eval = bool(self.is_decoder and not is_cross_attention and kv_len == N)

        context = None
        triton_ok = False
        if triton_dispatch is not None and not is_cross_attention:
            try:
                # Triton path expects FP16 inputs. Force FP16 only if values are bounded.
                # ELU+1 can produce values up to ~10, times cos/sin up to 1 → safe.
                Q_t = Q_cos.half()
                K_t = K_cos.half()
                K_s_t = K_sin.half()
                V_t = V_f.half()

                # Quick sanity check for overflow before calling Triton
                if (
                    torch.isnan(Q_t).any()
                    or torch.isnan(K_t).any()
                    or torch.isinf(Q_t).any()
                    or torch.isinf(K_t).any()
                ):
                    raise RuntimeError("Non-finite values detected before Triton dispatch")

                dispatch_kwargs = dict(self.triton_kwargs)
                dispatch_kwargs["cos_weight"] = cos_w.half()
                dispatch_kwargs["sin_weight"] = sin_w.half()
                dispatch_kwargs["eps"] = self.eps

                # Some dispatchers expect (Q_cos, Q_sin, K_cos, K_sin, V)
                context = triton_dispatch(
                    self.triton_kernel,
                    Q_t,
                    K_t,
                    V_t,
                    Q_sin=Q_k.half() * sin_w.half(),
                    K_sin=K_s_t,
                    is_causal=is_causal_eval,
                    **dispatch_kwargs,
                )
                context = context.float()
                triton_ok = True
            except Exception as e:
                logger.debug("Triton dispatch failed, using Python fallback: %s", e)
                triton_ok = False

        # ------------------------------------------------------------------
        # 7b. Python fallback: stable causal / bidirectional linear attention
        # ------------------------------------------------------------------
        if not triton_ok or context is None:
            if is_cross_attention:
                # Cross-attention: only bidirectional path (no causal masking)
                context = self._forward_bidirectional(
                    Q_cos,
                    Q_sin,
                    K_cos,
                    K_sin,
                    V_f,
                    B,
                    H,
                    d,
                    kv_len,
                    device,
                    torch.float32,
                )
            elif is_causal_eval:
                # Causal self-attention for decoders
                context = self._forward_causal(
                    Q_cos,
                    Q_sin,
                    K_cos,
                    K_sin,
                    V_f,
                    B,
                    H,
                    d,
                    N,
                    device,
                    torch.float32,
                )
            else:
                # Bidirectional self-attention for encoders
                context = self._forward_bidirectional(
                    Q_cos,
                    Q_sin,
                    K_cos,
                    K_sin,
                    V_f,
                    B,
                    H,
                    d,
                    N,
                    device,
                    torch.float32,
                )

        # ------------------------------------------------------------------
        # 8. Safety: kill any nan/inf before head mask and out_proj
        # ------------------------------------------------------------------
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # ------------------------------------------------------------------
        # 9. Head mask
        # ------------------------------------------------------------------
        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.shape[0], head_mask.shape[1], 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            context = context * head_mask

        # ------------------------------------------------------------------
        # 10. Merge heads → [B, N, H*d]
        # ------------------------------------------------------------------
        context = context.to(input_dtype)
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        # ------------------------------------------------------------------
        # 11. Output projection (or nested encoder bypass)
        # ------------------------------------------------------------------
        if self.is_nested_encoder:
            final_out = context
        else:
            final_out = self.out_proj(context)

        # ------------------------------------------------------------------
        # 12. Final safety clamp — prevents logit explosion downstream
        #     (adjust bounds if you know your task needs larger values)
        # ------------------------------------------------------------------
        final_out = torch.nan_to_num(final_out, nan=0.0, posinf=1e4, neginf=-1e4)
        final_out = final_out.clamp(min=-1e4, max=1e4)

        return final_out.to(input_dtype), None

    # ------------------------------------------------------------------
    # Causal forward (Fix 8: V masked in last chunk)
    # ------------------------------------------------------------------
    def _forward_causal(
        self,
        Q_cos: torch.Tensor,
        Q_sin: torch.Tensor,
        K_cos: torch.Tensor,
        K_sin: torch.Tensor,
        V: torch.Tensor,
        B: int,
        H: int,
        d: int,
        N: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Chunked causal linear attention. Stable in FP32."""
        chunk_size = max(16, min(64, N // 8))

        if chunk_size >= N:
            # Single-chunk: explicit causal mask
            causal_mask = torch.tril(torch.ones(N, N, device=device, dtype=torch.bool))
            score_c = torch.matmul(Q_cos, K_cos.transpose(-1, -2).contiguous())
            score_s = torch.matmul(Q_sin, K_sin.transpose(-1, -2).contiguous())
            scores = (score_c + score_s).masked_fill(~causal_mask, 0.0)
            num = torch.matmul(scores, V)
            denom = (score_c + score_s).masked_fill(~causal_mask, 0.0).sum(dim=-1, keepdim=True)
            return (num / denom.clamp(min=self.eps)).to(dtype)

        # Multi-chunk
        pad_len = (chunk_size - (N % chunk_size)) % chunk_size
        if pad_len > 0:
            Q_cos = F.pad(Q_cos, (0, 0, 0, pad_len))
            Q_sin = F.pad(Q_sin, (0, 0, 0, pad_len))
            K_cos = F.pad(K_cos, (0, 0, 0, pad_len))
            K_sin = F.pad(K_sin, (0, 0, 0, pad_len))
            V = F.pad(V, (0, 0, 0, pad_len))

        total_N = N + pad_len
        num_chunks = total_N // chunk_size

        qc_c = Q_cos.view(B, H, num_chunks, chunk_size, d)
        qs_c = Q_sin.view(B, H, num_chunks, chunk_size, d)
        kc_c = K_cos.view(B, H, num_chunks, chunk_size, d)
        ks_c = K_sin.view(B, H, num_chunks, chunk_size, d)
        v_c = V.view(B, H, num_chunks, chunk_size, d)

        causal_mask = torch.tril(
            torch.ones(chunk_size, chunk_size, device=device, dtype=torch.bool)
        )

        score_c = torch.matmul(qc_c, kc_c.transpose(-1, -2).contiguous()).masked_fill(
            ~causal_mask, 0.0
        )
        score_s = torch.matmul(qs_c, ks_c.transpose(-1, -2).contiguous()).masked_fill(
            ~causal_mask, 0.0
        )
        num_intra = torch.matmul(score_c + score_s, v_c)
        denom_intra = (score_c + score_s).sum(dim=-1, keepdim=True)

        S_cos = torch.zeros(B, H, d, d, device=device, dtype=torch.float32)
        S_sin = torch.zeros(B, H, d, d, device=device, dtype=torch.float32)
        Z_cos = torch.zeros(B, H, d, 1, device=device, dtype=torch.float32)
        Z_sin = torch.zeros(B, H, d, 1, device=device, dtype=torch.float32)

        num_list, denom_list = [], []
        for c in range(num_chunks):
            num_inter = torch.matmul(qc_c[:, :, c], S_cos) + torch.matmul(qs_c[:, :, c], S_sin)
            denom_inter = torch.matmul(qc_c[:, :, c], Z_cos) + torch.matmul(qs_c[:, :, c], Z_sin)

            num_list.append(num_intra[:, :, c] + num_inter)
            denom_list.append(denom_intra[:, :, c] + denom_inter)

            kc_curr = kc_c[:, :, c]
            ks_curr = ks_c[:, :, c]
            v_curr = v_c[:, :, c]

            if c == num_chunks - 1 and pad_len > 0:
                valid = chunk_size - pad_len
                kc_curr = kc_curr.clone()
                kc_curr[:, :, valid:] = 0.0
                ks_curr = ks_curr.clone()
                ks_curr[:, :, valid:] = 0.0
                v_curr = v_curr.clone()
                v_curr[:, :, valid:] = 0.0

            S_cos = S_cos + torch.matmul(kc_curr.transpose(-1, -2).contiguous(), v_curr)
            S_sin = S_sin + torch.matmul(ks_curr.transpose(-1, -2).contiguous(), v_curr)
            Z_cos = Z_cos + kc_curr.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous()
            Z_sin = Z_sin + ks_curr.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous()

        num = torch.cat(num_list, dim=2)[:, :, :N, :]
        denom = torch.cat(denom_list, dim=2)[:, :, :N, :]
        return (num / (denom + self.eps).clamp(min=self.eps)).to(dtype)

    # ------------------------------------------------------------------
    # Bidirectional forward
    # ------------------------------------------------------------------
    def _forward_bidirectional(
        self,
        Q_cos: torch.Tensor,
        Q_sin: torch.Tensor,
        K_cos: torch.Tensor,
        K_sin: torch.Tensor,
        V: torch.Tensor,
        B: int,
        H: int,
        d: int,
        N: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        KV_cos = torch.matmul(K_cos.transpose(-2, -1).contiguous(), V)
        KV_sin = torch.matmul(K_sin.transpose(-2, -1).contiguous(), V)

        Z_cos = K_cos.sum(dim=2, keepdim=True)  # [B, H, 1, d]
        Z_sin = K_sin.sum(dim=2, keepdim=True)

        num = torch.matmul(Q_cos, KV_cos) + torch.matmul(Q_sin, KV_sin)
        denom = torch.matmul(Q_cos, Z_cos.transpose(-2, -1).contiguous()) + torch.matmul(
            Q_sin, Z_sin.transpose(-2, -1).contiguous()
        )
        denom = denom.clamp(min=self.eps)

        context = (num / denom).to(dtype)
        return torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

Cosformer Attention with Linear Complexity and Exact Multi-Architecture Support.

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 triton_kernel

The type of the None singleton.

Instance variables

prop W_k : LoRALinear

Expand source code

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

prop W_q : LoRALinear

Expand source code

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

prop W_v : LoRALinear

Expand source code

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

prop c_proj : LoRALinear

Expand source code

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

prop dense : LoRALinear

Expand source code

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

prop k_lin : LoRALinear

Expand source code

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

prop key : LoRALinear

Expand source code

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

prop o_proj : LoRALinear

Expand source code

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

prop out_lin : LoRALinear

Expand source code

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

prop q_lin : LoRALinear

Expand source code

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

prop query : LoRALinear

Expand source code

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

prop v_lin : LoRALinear

Expand source code

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

prop value : LoRALinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    """Transfer pretrained projection weights with strict Frobenius norm verification."""
    if src is None:
        return

    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")
    )
    if cls_name in ["BertSelfAttention", "RobertaSelfAttention"] or (
        hasattr(src, "query") and not has_out
    ):
        self.is_nested_encoder = True
    if hasattr(src, "self") and hasattr(src, "output"):
        self.is_nested_encoder = True

    if hasattr(src, "rotary_emb"):
        self.rotary_emb = src.rotary_emb

    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 + c_proj
    if hasattr(src, "c_attn") and src.c_attn is not None:
        c_attn_mod = src.c_attn
        if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
            try:
                c_attn_mod.merge()
            except Exception as e:
                logger.debug(f"c_attn_mod.merge failed: {e}")
        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()
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or hasattr(c_attn_mod, "nf")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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()
                if (
                    (c_proj_mod.__class__.__name__ == "Conv1D")
                    or (base_o.__class__.__name__ == "Conv1D")
                    or hasattr(c_proj_mod, "nf")
                    or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                ):
                    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. RoBERTa / BERT nested self-attention
    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)
            base_mod = getattr(mod, "base_layer", mod)
            if base_mod is not None and hasattr(base_mod, "weight"):
                extracted[proj_key] = (
                    base_mod.weight.data.clone(),
                    (
                        base_mod.bias.data.clone()
                        if getattr(base_mod, "bias", None) is not None
                        else None
                    ),
                )
        if hasattr(src, "output") and hasattr(src.output, "dense"):
            out_mod = src.output.dense
            base_out = getattr(out_mod, "base_layer", out_mod)
            extracted["o"] = (
                base_out.weight.data.clone(),
                (
                    base_out.bias.data.clone()
                    if getattr(base_out, "bias", None) is not None
                    else None
                ),
            )

    # 3. Standard linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q": ["q_lin", "q_proj", "query", "W_q"],
            "k": ["k_lin", "k_proj", "key", "W_k"],
            "v": ["v_lin", "v_proj", "value", "W_v"],
            "o": ["out_lin", "out_proj", "dense", "o_proj", "c_proj", "wo"],
        }
        for proj_key, names in proj_candidates.items():
            for name in names:
                if hasattr(src, name):
                    mod = getattr(src, name)
                    base_mod = getattr(mod, "base_layer", mod)
                    if hasattr(base_mod, "weight"):
                        w_data = base_mod.weight.data.clone()
                        if (
                            (base_mod.__class__.__name__ == "Conv1D")
                            or hasattr(base_mod, "nf")
                            or (
                                w_data.ndim == 2
                                and proj_key == "o"
                                and w_data.shape[0] < w_data.shape[1]
                            )
                        ):
                            w_data = w_data.t().contiguous()
                        b_data = getattr(base_mod, "bias", None)
                        extracted[proj_key] = (
                            w_data,
                            b_data.data.clone() if b_data is not None else None,
                        )
                        break

    # Fallback via UniversalAdapter
    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(f"UniversalAdapter fallback: {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"CosformerAttention copy_weights: missing mandatory projections "
            f"{missing_after} from {src.__class__.__name__}."
        )

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device, dtype=target_mod.weight.dtype
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

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

            if self.freeze_base:
                target_mod.freeze_base_weights()

Transfer pretrained projection weights with strict Frobenius norm verification.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class DecayGatedLinearAttention (*args, **kwargs)

Expand source code

@ATTENTION_REGISTRY.register("svd_linear_decay")
@ATTENTION_REGISTRY.register("decay_gated_linear_attention")
class DecayGatedLinearAttention(SVDLinearAttention):
    r"""Decay Gated SVD Linear Attention."""

    is_custom_attention: bool = True

    def __init__(self, *args, **kwargs):
        kwargs["use_decay_gating"] = True
        super().__init__(*args, **kwargs)

Decay Gated SVD Linear Attention.

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

Ancestors

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

Inherited members

• SVDLinearAttention:
  • copy_weights
  • extra_repr
  • forward
  • is_custom_attention
  • supports_bidirectional
  • supports_causal
  • supports_gqa
  • triton_kernel

class DecayLinearAttention (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
num_heads: int | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "decay_linear_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "decay_linear_attn",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class DecayLinearAttention(BaseAttention):
    r"""Decay Linear Attention Module with Data-Dependent Exponential Decay Gating."""

    is_custom_attention: bool = True
    triton_kernel: str = "linear_attn"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        def _cfg_get(k: str, default: Any = None) -> Any:
            for obj in (config, comp_config):
                if obj is not None and hasattr(obj, k):
                    val = getattr(obj, k)
                    if val is not None:
                        return val
            return kwargs.get(k, default)

        hidden_size = kwargs.get("hidden_size")
        if hidden_size is None and d_model is not None:
            hidden_size = d_model
        if hidden_size is None:
            hidden_size = _cfg_get("hidden_size", _cfg_get("dim", None))
            if hidden_size is None and source_attn is not None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                hidden_size = (
                    getattr(out_layer, "out_features", 768) if out_layer is not None else 768
                )
            elif hidden_size is None:
                hidden_size = 768

        if hasattr(hidden_size, "hidden_size"):
            hidden_size = hidden_size.hidden_size

        if num_heads is None:
            num_heads = kwargs.get("num_attention_heads")
        if num_heads is None:
            num_heads = _cfg_get("num_attention_heads", _cfg_get("num_heads", None))
            if num_heads is None and source_attn is not None:
                num_heads = getattr(
                    source_attn,
                    "n_heads",
                    getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(source_attn, "num_heads", 12),
                    ),
                )
            elif num_heads is None:
                num_heads = 12

        super().__init__(config, rank)
        self.dim: int = int(hidden_size)
        self.num_heads: int = int(num_heads)

        resolved_head_dim = _cfg_get("head_dim", None)
        if resolved_head_dim is None and source_attn is not None:
            resolved_head_dim = getattr(source_attn, "head_dim", None)
        if (
            resolved_head_dim is None
            and source_attn is not None
            and hasattr(source_attn, "q_proj")
            and hasattr(source_attn.q_proj, "out_features")
        ):
            resolved_head_dim = source_attn.q_proj.out_features // max(1, self.num_heads)
        if resolved_head_dim is None:
            resolved_head_dim = self.dim // max(1, self.num_heads)
        self.head_dim: int = int(resolved_head_dim)

        kv_heads = _cfg_get("num_key_value_heads", None)
        if kv_heads is None and source_attn is not None:
            kv_heads = getattr(
                source_attn,
                "num_key_value_heads",
                getattr(
                    source_attn,
                    "num_kv_heads",
                    getattr(source_attn, "n_kv_heads", None),
                ),
            )
            if (
                kv_heads is None
                and hasattr(source_attn, "k_proj")
                and hasattr(source_attn.k_proj, "out_features")
                and self.head_dim > 0
            ):
                kv_heads = source_attn.k_proj.out_features // self.head_dim
        self.num_kv_heads: int = int(kv_heads if kv_heads is not None else self.num_heads)

        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 = max(1, self.num_heads // max(1, self.num_kv_heads))

        # Static layer index
        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx: int = int(resolved_layer_id)

        # LoRA adaptation configuration
        self.lora_rank: int = int(_cfg_get("lora_rank", kwargs.get("lora_rank", 16)))
        self.lora_alpha: float = float(_cfg_get("lora_alpha", kwargs.get("lora_alpha", 32.0)))
        self.freeze_base: bool = self.lora_rank > 0

        # Decoder & nested encoder detection
        model_type = getattr(config, "model_type", "").lower() if config is not None else ""
        is_dec = _cfg_get("is_decoder", _cfg_get("is_causal", kwargs.get("is_causal", False)))
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "phi",
            "gemma",
            "opt",
            "bloom",
            "falcon",
        )
        if (
            any(model_type.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)
        self.is_nested_encoder: bool = bool(kwargs.get("is_nested_encoder", False))

        # Multi-head projections equipped with LoRA adapters
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.decay_gate = nn.Linear(self.dim, self.num_heads, bias=True)
        # Bounded initialization spanning [sigmoid(-0.5), sigmoid(0.0)] = [0.38, 0.50] for head diversity and wide context horizon
        self.log_decay = nn.Parameter(torch.linspace(2.0, 4.0, self.num_heads))
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        nn.init.zeros_(self.decay_gate.weight)
        nn.init.zeros_(self.decay_gate.bias)

        self.rotary_emb = None
        self.comp_config: Any | None = comp_config
        self.triton_kwargs = {"phi": "elu+1", "eps": 1e-6}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    def copy_weights(self, old_attn: Any) -> None:
        """Transfer pretrained projection weights with strict Frobenius norm verification."""
        if old_attn is None:
            return

        cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        has_out_lin = (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
            or hasattr(old_attn, "dense")
        )
        if (
            (
                hasattr(old_attn, "self")
                and hasattr(old_attn.self, "query")
                and hasattr(old_attn, "output")
                and hasattr(old_attn.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
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "rotary_emb"):
            self.rotary_emb = old_attn.rotary_emb

        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 + c_proj
        if hasattr(old_attn, "c_attn") and old_attn.c_attn is not None:
            c_attn_mod = old_attn.c_attn
            if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    logger.debug(f"c_attn_mod.merge failed: {e}")
            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()
                if (
                    ("Conv1D" in c_attn_mod.__class__.__name__)
                    or ("Conv1D" in base_c.__class__.__name__)
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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(old_attn, "c_proj") and old_attn.c_proj is not None:
                c_proj_mod = old_attn.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()
                    if (
                        ("Conv1D" in c_proj_mod.__class__.__name__)
                        or ("Conv1D" in base_o.__class__.__name__)
                        or hasattr(c_proj_mod, "nf")
                        or hasattr(base_o, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        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. RoBERTa / BERT nested self-attention
        elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
            s = old_attn.self
            for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]:
                mod = getattr(s, sub_name, None)
                base_mod = getattr(mod, "base_layer", mod)
                if base_mod is not None and hasattr(base_mod, "weight"):
                    extracted[proj_key] = (
                        base_mod.weight.data.clone(),
                        (
                            base_mod.bias.data.clone()
                            if getattr(base_mod, "bias", None) is not None
                            else None
                        ),
                    )
            if hasattr(old_attn, "output") and hasattr(old_attn.output, "dense"):
                out_mod = old_attn.output.dense
                base_out = getattr(out_mod, "base_layer", out_mod)
                extracted["o"] = (
                    base_out.weight.data.clone(),
                    (
                        base_out.bias.data.clone()
                        if getattr(base_out, "bias", None) is not None
                        else None
                    ),
                )

        # 3. Standard linear projections (DistilBERT, Qwen, LLaMA)
        else:
            proj_candidates = {
                "q": ["q_lin", "q_proj", "query", "W_q"],
                "k": ["k_lin", "k_proj", "key", "W_k"],
                "v": ["v_lin", "v_proj", "value", "W_v"],
                "o": ["out_lin", "out_proj", "dense", "o_proj", "c_proj", "wo"],
            }
            for proj_key, names in proj_candidates.items():
                for name in names:
                    if hasattr(old_attn, name):
                        mod = getattr(old_attn, name)
                        base_mod = getattr(mod, "base_layer", mod)
                        if hasattr(base_mod, "weight"):
                            w_data = base_mod.weight.data.clone()
                            if (
                                (base_mod.__class__.__name__ == "Conv1D")
                                or hasattr(base_mod, "nf")
                                or (
                                    w_data.ndim == 2
                                    and proj_key == "o"
                                    and w_data.shape[0] < w_data.shape[1]
                                )
                            ):
                                w_data = w_data.t().contiguous()
                            b_data = getattr(base_mod, "bias", None)
                            extracted[proj_key] = (
                                w_data,
                                b_data.data.clone() if b_data is not None else None,
                            )
                            break

        # Fallback via UniversalAdapter
        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(old_attn, 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(f"UniversalAdapter fallback: {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"DecayLinearAttention copy_weights: missing mandatory projections {missing_after} "
                f"from {old_attn.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

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

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    def _decay_scan_linear(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        lambdas: torch.Tensor,
        mask: torch.Tensor | None = None,
        reverse: bool = False,
        chunk_size: int = 32,
    ) -> torch.Tensor:
        r"""
        Numerically stable parallel decay scan with O(1) state per chunk (chunked recurrence).
        Eliminates quadratic O(L^2) memory footprint and prevents GPU out-of-memory.
        """
        B, H, L, d = q.shape

        # Path 1: Hardware-fused Triton kernel on CUDA
        if (
            HAS_TRITON_DECAY
            and q.is_cuda
            and d <= MAX_REGISTER_DIM
            and mask is None
            and not reverse
        ):
            try:
                return gla_scan_triton(q.float(), k.float(), v.float(), lambdas.float()).to(q.dtype)
            except Exception as e:
                logger.debug("Triton decay scan failed, falling back to chunked PyTorch: %s", e)

        q_f = q.float()
        k_f = k.float()
        v_f = v.float()
        lambdas_f = lambdas.float().clamp(min=1e-4, max=1.0 - 1e-4)

        if mask is not None:
            # mask is [B, L] boolean tensor where True = valid token
            m_f = mask.float().unsqueeze(1).unsqueeze(-1)  # [B, 1, L, 1]
            k_f = k_f * m_f
            v_f = v_f * m_f

        if reverse:
            q_f = q_f.flip(dims=[2])
            k_f = k_f.flip(dims=[2])
            v_f = v_f.flip(dims=[2])
            lambdas_f = lambdas_f.flip(dims=[-1])

        log_lambdas = torch.log(lambdas_f)
        cum_log = torch.cumsum(log_lambdas, dim=-1)

        if chunk_size >= L:
            log_decay = cum_log.unsqueeze(-1) - cum_log.unsqueeze(-2)
            decay = torch.exp(log_decay.clamp(max=0.0))
            causal_mask = torch.tril(torch.ones(L, L, device=q.device, dtype=torch.bool))
            decay = decay.masked_fill(~causal_mask, 0.0)
            scores = torch.matmul(q_f, k_f.transpose(-1, -2).contiguous()) * decay
            num = torch.matmul(scores, v_f)
            denom = scores.sum(dim=-1, keepdim=True)
        else:
            pad_len = (chunk_size - (L % chunk_size)) % chunk_size
            if pad_len > 0:
                q_pad = F.pad(q_f, (0, 0, 0, pad_len))
                k_pad = F.pad(k_f, (0, 0, 0, pad_len))
                v_pad = F.pad(v_f, (0, 0, 0, pad_len))
                last_cum = cum_log[:, :, -1:]
                cum_pad = torch.cat([cum_log, last_cum.expand(-1, -1, pad_len)], dim=-1)
            else:
                q_pad, k_pad, v_pad, cum_pad = q_f, k_f, v_f, cum_log

            total_L = L + pad_len
            num_chunks = total_L // chunk_size

            qc = q_pad.view(B, H, num_chunks, chunk_size, d)
            kc = k_pad.view(B, H, num_chunks, chunk_size, d)
            vc = v_pad.view(B, H, num_chunks, chunk_size, d)
            cum_c = cum_pad.view(B, H, num_chunks, chunk_size)

            intra_log_decay = cum_c.unsqueeze(-1) - cum_c.unsqueeze(-2)
            intra_decay = torch.exp(intra_log_decay.clamp(max=0.0))
            causal_mask = torch.tril(
                torch.ones(chunk_size, chunk_size, device=q.device, dtype=torch.bool)
            )
            intra_decay = intra_decay.masked_fill(~causal_mask, 0.0)

            if pad_len > 0:
                valid_in_last = chunk_size - pad_len
                intra_decay[:, :, -1, :, valid_in_last:] = 0.0

            intra_scores = torch.matmul(qc, kc.transpose(-1, -2).contiguous()) * intra_decay
            num_intra = torch.matmul(intra_scores, vc)
            denom_intra = intra_scores.sum(dim=-1, keepdim=True)

            S = torch.zeros(B, H, d, d, device=q.device, dtype=torch.float32)
            z = torch.zeros(B, H, d, 1, device=q.device, dtype=torch.float32)

            num_list = []
            denom_list = []

            prev_cum_end = None
            for c in range(num_chunks):
                curr_cum = cum_c[:, :, c]
                if c == 0:
                    inter_decay = torch.zeros(
                        B, H, chunk_size, 1, device=q.device, dtype=torch.float32
                    )
                else:
                    inter_decay = torch.exp((curr_cum - prev_cum_end).clamp(max=0.0)).unsqueeze(-1)

                num_inter = torch.matmul(qc[:, :, c], S) * inter_decay
                denom_inter = torch.matmul(qc[:, :, c], z) * inter_decay

                num_list.append(num_intra[:, :, c] + num_inter)
                denom_list.append(denom_intra[:, :, c] + denom_inter)

                cum_end = curr_cum[:, :, -1:]
                decay_to_end = torch.exp((cum_end - curr_cum).clamp(max=0.0)).unsqueeze(-1)

                k_curr = kc[:, :, c]
                v_curr = vc[:, :, c]
                if c == num_chunks - 1 and pad_len > 0:
                    valid_in_last = chunk_size - pad_len
                    k_curr = k_curr.clone()
                    k_curr[:, :, valid_in_last:] = 0.0

                k_to_end = k_curr * decay_to_end

                if c == 0:
                    decay_prev_to_end = 0.0
                else:
                    decay_prev_to_end = torch.exp(
                        (cum_end - prev_cum_end).clamp(max=0.0)
                    ).unsqueeze(-1)

                S = S * decay_prev_to_end + torch.matmul(
                    k_to_end.transpose(-1, -2).contiguous(), v_curr
                )
                z = (
                    z * decay_prev_to_end
                    + k_to_end.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous()
                )
                prev_cum_end = cum_end

            num = torch.cat(num_list, dim=2)[:, :, :L, :]
            denom = torch.cat(denom_list, dim=2)[:, :, :L, :]

        y = num / denom.clamp(min=1e-6)

        if reverse:
            y = y.flip(dims=[2])

        return y.to(q.dtype)

    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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim

        is_cross_attention = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross_attention else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Projections
        q = self.q_proj(hidden_states).view(B, L, H, d).transpose(1, 2).contiguous()
        k = self.k_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()
        v = self.v_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # 2. Rotary Position Embeddings (RoPE)
        if position_embeddings is not None and not is_cross_attention:
            if isinstance(position_embeddings, tuple):
                cos, sin = position_embeddings
            else:
                cos, sin = position_embeddings, None
            if sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        elif self.rotary_emb is not None and not is_cross_attention:
            if position_ids is not None:
                cos, sin = self.rotary_emb(v, seq_len=kv_len)
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)

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

        # 4. Kernel feature mapping
        q = F.elu(q) + 1.0
        k = F.elu(k) + 1.0

        # Data-dependent decay lambdas in (0, 1) with bounded negative log decay bias
        decay_logits = self.decay_gate(hidden_states).transpose(
            1, 2
        ).contiguous() + self.log_decay.view(1, -1, 1)
        lambdas = torch.sigmoid(decay_logits).clamp(min=0.85, max=0.999)

        eff_mask = encoder_attention_mask if is_cross_attention else attention_mask
        mask_bool = None
        if eff_mask is not None:
            mask_bool = get_token_valid_mask(eff_mask, B, kv_len, hidden_states.device)

        # 5. True O(L) Linear Decay Scan
        y_fwd = self._decay_scan_linear(q, k, v, lambdas, mask=mask_bool, reverse=False)

        if self.is_decoder or is_cross_attention:
            y = y_fwd
        else:
            y_bwd = self._decay_scan_linear(q, k, v, lambdas, mask=mask_bool, reverse=True)
            y = 0.5 * (y_fwd + y_bwd)

        # 6. Head mask handling
        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.shape[0], head_mask.shape[1], 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            y = y * head_mask

        y = torch.nan_to_num(y, nan=0.0)
        out = y.transpose(1, 2).contiguous().view(B, L, self.all_head_size)

        # 7. Output projection with nested encoder bypass (no output_scale)
        if self.is_nested_encoder:
            final_out = out
        else:
            final_out = self.out_proj(out)

        if mask_bool is not None and not is_cross_attention:
            final_out = final_out * mask_bool.unsqueeze(-1).to(final_out.dtype)

        return final_out, None

Decay Linear Attention Module with Data-Dependent Exponential Decay Gating.

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 triton_kernel : str

The type of the None singleton.

Instance variables

prop W_k : LoRALinear

Expand source code

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

prop W_q : LoRALinear

Expand source code

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

prop W_v : LoRALinear

Expand source code

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

prop c_proj : LoRALinear

Expand source code

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

prop dense : LoRALinear

Expand source code

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

prop k_lin : LoRALinear

Expand source code

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

prop key : LoRALinear

Expand source code

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

prop o_proj : LoRALinear

Expand source code

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

prop out_lin : LoRALinear

Expand source code

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

prop q_lin : LoRALinear

Expand source code

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

prop query : LoRALinear

Expand source code

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

prop v_lin : LoRALinear

Expand source code

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

prop value : LoRALinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Transfer pretrained projection weights with strict Frobenius norm verification."""
    if old_attn is None:
        return

    cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
    has_out_lin = (
        hasattr(old_attn, "out_lin")
        or hasattr(old_attn, "c_proj")
        or hasattr(old_attn, "o_proj")
        or hasattr(old_attn, "dense")
    )
    if (
        (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.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
    ):
        self.is_nested_encoder = True

    if hasattr(old_attn, "rotary_emb"):
        self.rotary_emb = old_attn.rotary_emb

    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 + c_proj
    if hasattr(old_attn, "c_attn") and old_attn.c_attn is not None:
        c_attn_mod = old_attn.c_attn
        if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
            try:
                c_attn_mod.merge()
            except Exception as e:
                logger.debug(f"c_attn_mod.merge failed: {e}")
        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()
            if (
                ("Conv1D" in c_attn_mod.__class__.__name__)
                or ("Conv1D" in base_c.__class__.__name__)
                or hasattr(c_attn_mod, "nf")
                or hasattr(base_c, "nf")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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(old_attn, "c_proj") and old_attn.c_proj is not None:
            c_proj_mod = old_attn.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()
                if (
                    ("Conv1D" in c_proj_mod.__class__.__name__)
                    or ("Conv1D" in base_o.__class__.__name__)
                    or hasattr(c_proj_mod, "nf")
                    or hasattr(base_o, "nf")
                    or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                ):
                    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. RoBERTa / BERT nested self-attention
    elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
        s = old_attn.self
        for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]:
            mod = getattr(s, sub_name, None)
            base_mod = getattr(mod, "base_layer", mod)
            if base_mod is not None and hasattr(base_mod, "weight"):
                extracted[proj_key] = (
                    base_mod.weight.data.clone(),
                    (
                        base_mod.bias.data.clone()
                        if getattr(base_mod, "bias", None) is not None
                        else None
                    ),
                )
        if hasattr(old_attn, "output") and hasattr(old_attn.output, "dense"):
            out_mod = old_attn.output.dense
            base_out = getattr(out_mod, "base_layer", out_mod)
            extracted["o"] = (
                base_out.weight.data.clone(),
                (
                    base_out.bias.data.clone()
                    if getattr(base_out, "bias", None) is not None
                    else None
                ),
            )

    # 3. Standard linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q": ["q_lin", "q_proj", "query", "W_q"],
            "k": ["k_lin", "k_proj", "key", "W_k"],
            "v": ["v_lin", "v_proj", "value", "W_v"],
            "o": ["out_lin", "out_proj", "dense", "o_proj", "c_proj", "wo"],
        }
        for proj_key, names in proj_candidates.items():
            for name in names:
                if hasattr(old_attn, name):
                    mod = getattr(old_attn, name)
                    base_mod = getattr(mod, "base_layer", mod)
                    if hasattr(base_mod, "weight"):
                        w_data = base_mod.weight.data.clone()
                        if (
                            (base_mod.__class__.__name__ == "Conv1D")
                            or hasattr(base_mod, "nf")
                            or (
                                w_data.ndim == 2
                                and proj_key == "o"
                                and w_data.shape[0] < w_data.shape[1]
                            )
                        ):
                            w_data = w_data.t().contiguous()
                        b_data = getattr(base_mod, "bias", None)
                        extracted[proj_key] = (
                            w_data,
                            b_data.data.clone() if b_data is not None else None,
                        )
                        break

    # Fallback via UniversalAdapter
    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(old_attn, 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(f"UniversalAdapter fallback: {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"DecayLinearAttention copy_weights: missing mandatory projections {missing_after} "
            f"from {old_attn.__class__.__name__}."
        )

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device, dtype=target_mod.weight.dtype
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

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

            if self.freeze_base:
                target_mod.freeze_base_weights()

Transfer pretrained projection weights with strict Frobenius norm verification.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class DualStageStripeAttention (cfg: NEXUSConfig,
head_dim: int,
num_heads: int = 1,
is_decoder: bool = False,
is_nested_encoder: bool = False,
**kwargs)

Expand source code

class DualStageStripeAttention(nn.Module):
    def __init__(
        self,
        cfg: NEXUSConfig,
        head_dim: int,
        num_heads: int = 1,
        is_decoder: bool = False,
        is_nested_encoder: bool = False,
        **kwargs,
    ) -> None:
        super().__init__()
        self.cfg = cfg
        self.head_dim = head_dim
        self.num_heads = num_heads
        self.is_decoder = is_decoder
        self.is_nested_encoder = is_nested_encoder
        self.xattn_scorer = XAttentionScorer()

    def _inter_block_topk(self, scores: torch.Tensor, top_k: int) -> torch.Tensor:
        B, H, nb, nb2 = scores.shape
        tk = min(top_k, nb2)
        scores_inter = scores.clone()
        min_dim = min(nb, nb2)
        diag_idx = torch.arange(min_dim, device=scores.device)
        scores_inter[:, :, diag_idx, diag_idx] = float("-inf")
        _, idx = scores_inter.topk(tk, dim=-1)
        mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, idx, True)
        mask[:, :, diag_idx, diag_idx] = True
        return mask

    def forward(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        token_valid: torch.Tensor | None = None,
        seer_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, N, D = Q.shape
        bs = self.cfg.block_size
        nb = (N + bs - 1) // bs
        top_k = self.cfg.top_k_blocks
        scale = D**-0.5

        if self.cfg.xattn_enabled:
            block_scores = self.xattn_scorer.score_blocks(Q, K, bs, self.cfg.xattn_stride)
            inter_mask = self._inter_block_topk(block_scores, top_k)
        else:
            inter_mask = torch.ones(B, H, nb, nb, device=Q.device, dtype=torch.bool)

        if seer_mask is not None:
            inter_mask = inter_mask | seer_mask

        inter_mask[:, :, :, 0] = True
        diag_idx = torch.arange(nb, device=Q.device)
        inter_mask[:, :, diag_idx, diag_idx] = True

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        expanded_mask = inter_mask.repeat_interleave(bs, dim=2).repeat_interleave(bs, dim=3)[
            :, :, :N, :N
        ]
        scores = scores.masked_fill(~expanded_mask, -10000.0)

        if self.cfg.stripe_enabled and N > 256:
            sw = self.cfg.stripe_width
            idx = torch.arange(N, device=Q.device)
            stripe_group = idx // sw
            q_stripe = stripe_group.unsqueeze(1)
            k_stripe = stripe_group.unsqueeze(0)
            stripe_mask = (q_stripe - k_stripe).abs() <= 1
            scores = scores.masked_fill(~stripe_mask.unsqueeze(0).unsqueeze(0), -10000.0)

        if self.is_decoder:
            causal = torch.tril(torch.ones(N, N, device=Q.device, dtype=torch.bool))
            scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -10000.0)

        if token_valid is not None:
            scores = scores.masked_fill((~token_valid.bool()).unsqueeze(1).unsqueeze(2), -10000.0)

        if self.cfg.softpick_enabled:
            attn = softpick(scores, dim=-1, temperature=self.cfg.softpick_temperature).to(Q.dtype)
        else:
            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            attn = F.softmax(scores, dim=-1).to(Q.dtype)

        context = torch.matmul(attn, V)
        return context

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self,
Q: torch.Tensor,
K: torch.Tensor,
V: torch.Tensor,
token_valid: torch.Tensor | None = None,
seer_mask: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def forward(
    self,
    Q: torch.Tensor,
    K: torch.Tensor,
    V: torch.Tensor,
    token_valid: torch.Tensor | None = None,
    seer_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    B, H, N, D = Q.shape
    bs = self.cfg.block_size
    nb = (N + bs - 1) // bs
    top_k = self.cfg.top_k_blocks
    scale = D**-0.5

    if self.cfg.xattn_enabled:
        block_scores = self.xattn_scorer.score_blocks(Q, K, bs, self.cfg.xattn_stride)
        inter_mask = self._inter_block_topk(block_scores, top_k)
    else:
        inter_mask = torch.ones(B, H, nb, nb, device=Q.device, dtype=torch.bool)

    if seer_mask is not None:
        inter_mask = inter_mask | seer_mask

    inter_mask[:, :, :, 0] = True
    diag_idx = torch.arange(nb, device=Q.device)
    inter_mask[:, :, diag_idx, diag_idx] = True

    scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

    expanded_mask = inter_mask.repeat_interleave(bs, dim=2).repeat_interleave(bs, dim=3)[
        :, :, :N, :N
    ]
    scores = scores.masked_fill(~expanded_mask, -10000.0)

    if self.cfg.stripe_enabled and N > 256:
        sw = self.cfg.stripe_width
        idx = torch.arange(N, device=Q.device)
        stripe_group = idx // sw
        q_stripe = stripe_group.unsqueeze(1)
        k_stripe = stripe_group.unsqueeze(0)
        stripe_mask = (q_stripe - k_stripe).abs() <= 1
        scores = scores.masked_fill(~stripe_mask.unsqueeze(0).unsqueeze(0), -10000.0)

    if self.is_decoder:
        causal = torch.tril(torch.ones(N, N, device=Q.device, dtype=torch.bool))
        scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -10000.0)

    if token_valid is not None:
        scores = scores.masked_fill((~token_valid.bool()).unsqueeze(1).unsqueeze(2), -10000.0)

    if self.cfg.softpick_enabled:
        attn = softpick(scores, dim=-1, temperature=self.cfg.softpick_temperature).to(Q.dtype)
    else:
        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(Q.dtype)

    context = torch.matmul(attn, V)
    return context

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class DynamicLowRankProjector (rank: int = 32,
in_dim: int = 768,
out_dim: int = 768,
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 DynamicLowRankProjector(nn.Module):
    is_custom_attention: bool = True
    r"""SVD-initialized low-rank projector with dynamic variational gate scaling and LoRA adapters."""

    def __init__(
        self,
        rank: int = 32,
        in_dim: int = 768,
        out_dim: int = 768,
        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.rank = rank
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.lora_rank = lora_rank
        self.scaling = lora_alpha / max(1, lora_rank)

        self.U = nn.Parameter(torch.randn(out_dim, rank, **factory_kwargs) * 0.02)
        self.S = nn.Parameter(torch.ones(rank, **factory_kwargs))
        self.V = nn.Parameter(torch.randn(in_dim, rank, **factory_kwargs) * 0.02)

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_dim, **factory_kwargs))
        else:
            self.register_parameter("bias", None)

        if lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty((lora_rank, in_dim), **factory_kwargs))
            self.lora_B = nn.Parameter(torch.empty((out_dim, lora_rank), **factory_kwargs))
            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)

        if freeze_base:
            self.freeze_base_weights()

    def freeze_base_weights(self) -> None:
        self.U.requires_grad = False
        self.S.requires_grad = False
        self.V.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 set_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
        with torch.no_grad():
            W = weight.float()
            U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            eff_rank = min(self.rank, U.size(1), Vh.size(0))

            self.U.data.zero_()
            self.S.data.zero_()
            self.V.data.zero_()

            min_u = min(self.U.size(0), U.size(0))
            min_v = min(self.V.size(0), Vh.size(1))

            self.U.data[:min_u, :eff_rank].copy_(U[:min_u, :eff_rank])
            self.S.data[:eff_rank].copy_(S[:eff_rank])
            self.V.data[:min_v, :eff_rank].copy_(Vh[:eff_rank, :min_v].t().contiguous())

            if bias is not None and self.bias is not None:
                b_val = bias.data if isinstance(bias, nn.Parameter) else bias
                min_b = min(self.bias.size(0), b_val.size(0))
                self.bias.data[:min_b].copy_(b_val[:min_b])

            rec = (self.U[:, :eff_rank] * self.S[:eff_rank].unsqueeze(0)) @ self.V[
                :, :eff_rank
            ].t().contiguous()
            f_norm = float(torch.linalg.norm(rec))
            if f_norm < 1e-4:
                logger.warning(
                    f"DynamicLowRankProjector: reconstructed norm is near zero ({f_norm}), re-initializing randomly"
                )
                nn.init.orthogonal_(self.U.data)
                nn.init.orthogonal_(self.V.data)
                self.S.data.fill_(1.0)

    def copy_lora_from(self, sub: Any) -> None:
        """Preserve source LoRA adapters."""
        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 self.lora_A is not None:
                r_min = min(self.lora_A.shape[0], src_lA.shape[0])
                in_min = min(self.lora_A.shape[1], src_lA.shape[1])
                self.lora_A.data[:r_min, :in_min].copy_(src_lA.data[:r_min, :in_min])
                self.lora_A.requires_grad = True
            if hasattr(src_lB, "data") and self.lora_B is not None:
                out_min = min(self.lora_B.shape[0], src_lB.shape[0])
                r_min = min(self.lora_B.shape[1], src_lB.shape[1])
                self.lora_B.data[:out_min, :r_min].copy_(src_lB.data[:out_min, :r_min])
                self.lora_B.requires_grad = True
            if hasattr(sub, "lora_scaling"):
                self.scaling = sub.lora_scaling

    def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
        v = self.V.to(x.dtype)
        s = self.S.to(x.dtype)
        u = self.U.to(x.dtype)
        g = gate.to(x.dtype)
        z = torch.matmul(x, v)
        z = z * s.unsqueeze(0).unsqueeze(0) * g.unsqueeze(1)
        out = torch.matmul(z, u.t().contiguous())
        if self.bias is not None:
            out = out + self.bias.to(x.dtype)

        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            la = self.lora_A.to(x.dtype)
            lb = self.lora_B.to(x.dtype)
            lora_out = F.linear(F.linear(x, la), lb) * self.scaling
            out = out + lora_out

        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

SVD-initialized low-rank projector with dynamic variational gate scaling and LoRA adapters.

Methods

def copy_lora_from(self,
sub: Any) ‑> None

Expand source code

def copy_lora_from(self, sub: Any) -> None:
    """Preserve source LoRA adapters."""
    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 self.lora_A is not None:
            r_min = min(self.lora_A.shape[0], src_lA.shape[0])
            in_min = min(self.lora_A.shape[1], src_lA.shape[1])
            self.lora_A.data[:r_min, :in_min].copy_(src_lA.data[:r_min, :in_min])
            self.lora_A.requires_grad = True
        if hasattr(src_lB, "data") and self.lora_B is not None:
            out_min = min(self.lora_B.shape[0], src_lB.shape[0])
            r_min = min(self.lora_B.shape[1], src_lB.shape[1])
            self.lora_B.data[:out_min, :r_min].copy_(src_lB.data[:out_min, :r_min])
            self.lora_B.requires_grad = True
        if hasattr(sub, "lora_scaling"):
            self.scaling = sub.lora_scaling

Preserve source LoRA adapters.

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

Expand source code

def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
    v = self.V.to(x.dtype)
    s = self.S.to(x.dtype)
    u = self.U.to(x.dtype)
    g = gate.to(x.dtype)
    z = torch.matmul(x, v)
    z = z * s.unsqueeze(0).unsqueeze(0) * g.unsqueeze(1)
    out = torch.matmul(z, u.t().contiguous())
    if self.bias is not None:
        out = out + self.bias.to(x.dtype)

    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        la = self.lora_A.to(x.dtype)
        lb = self.lora_B.to(x.dtype)
        lora_out = F.linear(F.linear(x, la), lb) * self.scaling
        out = out + lora_out

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    self.U.requires_grad = False
    self.S.requires_grad = False
    self.V.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 set_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) ‑> None

Expand source code

def set_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
    with torch.no_grad():
        W = weight.float()
        U, S, Vh = torch.linalg.svd(W, full_matrices=False)
        eff_rank = min(self.rank, U.size(1), Vh.size(0))

        self.U.data.zero_()
        self.S.data.zero_()
        self.V.data.zero_()

        min_u = min(self.U.size(0), U.size(0))
        min_v = min(self.V.size(0), Vh.size(1))

        self.U.data[:min_u, :eff_rank].copy_(U[:min_u, :eff_rank])
        self.S.data[:eff_rank].copy_(S[:eff_rank])
        self.V.data[:min_v, :eff_rank].copy_(Vh[:eff_rank, :min_v].t().contiguous())

        if bias is not None and self.bias is not None:
            b_val = bias.data if isinstance(bias, nn.Parameter) else bias
            min_b = min(self.bias.size(0), b_val.size(0))
            self.bias.data[:min_b].copy_(b_val[:min_b])

        rec = (self.U[:, :eff_rank] * self.S[:eff_rank].unsqueeze(0)) @ self.V[
            :, :eff_rank
        ].t().contiguous()
        f_norm = float(torch.linalg.norm(rec))
        if f_norm < 1e-4:
            logger.warning(
                f"DynamicLowRankProjector: reconstructed norm is near zero ({f_norm}), re-initializing randomly"
            )
            nn.init.orthogonal_(self.U.data)
            nn.init.orthogonal_(self.V.data)
            self.S.data.fill_(1.0)

class EATBlock (config: Any | None = None,
rank: int | None = 64,
d_model: int | Any = 768,
n_heads: int = 12,
prune_layer: Any | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "eat", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "eat_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class EATBlock(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        d_model: int | Any = 768,
        n_heads: int = 12,
        prune_layer: Any | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # ---- EAT-FIX-10: config None-safety ----
        cfg_obj = getattr(config, "text_config", config) if config is not None else None
        hidden_size = getattr(cfg_obj, "hidden_size", None) if cfg_obj is not None else None
        num_heads = getattr(cfg_obj, "num_attention_heads", None) if cfg_obj is not None else None

        if hidden_size is None and source_attn is not None:
            for cand in ("q_lin", "query", "q_proj", "c_attn", "qkv"):
                mod = getattr(source_attn, cand, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    hidden_size = w.shape[1] if w.ndim == 2 else None
                    break
                elif isinstance(mod, torch.Tensor) and mod.ndim == 2:
                    hidden_size = mod.shape[1]
                    break
        if hidden_size is None:
            hidden_size = kwargs.get(
                "hidden_size",
                kwargs.get("d_model", d_model if isinstance(d_model, int) else 768),
            )
        if num_heads is None:
            num_heads = kwargs.get(
                "num_attention_heads",
                kwargs.get("num_heads", kwargs.get("n_heads", n_heads)),
            )

        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        self.d_model = int(hidden_size)
        self.dim = self.d_model
        self.n_heads = int(num_heads)
        self.num_heads = self.n_heads

        # ---- num_kv_heads ----
        cfg_kv = None
        if config is not None:
            cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        if cfg_kv is None:
            cfg_kv = kwargs.get("num_kv_heads", self.num_heads)
        self.num_kv_heads = max(1, min(int(cfg_kv), self.num_heads))
        self.kv_repeat = max(1, self.num_heads // self.num_kv_heads)

        # ---- head_dim ----
        cfg_hd = getattr(config, "head_dim", None) if config else None
        self.head_dim = int(
            cfg_hd or kwargs.get("head_dim") or (self.dim // max(1, self.num_heads))
        )

        self.all_head_size = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim

        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx = int(resolved_layer_id)

        # ---- LoRA ----
        self.lora_rank = int(
            getattr(
                config,
                "lora_rank",
                (
                    getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16))
                    if comp_config is not None
                    else kwargs.get("lora_rank", 16)
                ),
            )
        )
        self.lora_alpha = float(
            getattr(
                config,
                "lora_alpha",
                (
                    getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
                    if comp_config is not None
                    else kwargs.get("lora_alpha", 32.0)
                ),
            )
        )
        self.freeze_base = self.lora_rank > 0

        # ---- decoder detection ----
        model_type = str(getattr(config, "model_type", "") if config else "").lower()
        is_dec = False
        if config is not None:
            is_dec = bool(
                getattr(config, "is_decoder", False) or getattr(config, "is_causal", False)
            )
        if any(
            model_type.startswith(p)
            for p in (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        if source_attn is not None:
            if getattr(source_attn, "is_decoder", False) or hasattr(source_attn, "c_attn"):
                is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = self.is_decoder

        # EAT-FIX-03: never mark DistilBERT-like as nested
        self.is_nested_encoder = False

        # ---- projections ----
        self.q_lin = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_lin = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_lin = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_lin = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        dropout_p = 0.1
        if config is not None:
            dropout_p = float(
                getattr(
                    config,
                    "dropout",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                )
            )
        elif comp_config is not None:
            dropout_p = float(getattr(comp_config, "dropout", 0.1))
        self.dropout = nn.Dropout(dropout_p)

        # ---- EAT-FIX-07: pruning only when explicitly requested ----
        self.prune_layer: Any | None = None
        if prune_layer is not None:
            self.prune_layer = prune_layer
        else:
            enable_pruning = False
            keep_ratio = None
            if comp_config is not None:
                enable_pruning = bool(getattr(comp_config, "enable_pruning", False))
                keep_ratio = getattr(comp_config, "keep_ratio", None)
            if kwargs.get("enable_pruning", False):
                enable_pruning = True
            if kwargs.get("keep_ratio") is not None:
                keep_ratio = kwargs["keep_ratio"]

            if enable_pruning and ProgressiveTokenPruning is not None:
                effective_ratio = float(keep_ratio) if keep_ratio is not None else 0.7
                try:
                    self.prune_layer = ProgressiveTokenPruning(
                        config=config,
                        comp_config=comp_config,
                        keep_ratio=effective_ratio,
                        layer_id=self.layer_idx,
                        layer_idx=self.layer_idx,
                    )
                except Exception as e:
                    logger.debug("EAT: ProgressiveTokenPruning not created: %s", e)
                    self.prune_layer = None

        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )
        self.comp_config = comp_config

        # EAT-FIX-09
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

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

    # ---- aliases ----
    @property
    def query(self):
        return self.q_lin

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

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

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

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

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

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

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

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

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

    # ------------------------------------------------------------------
    # copy_weights (EAT-FIX-11)
    # ------------------------------------------------------------------
    def copy_weights(self, pretrained_attn: nn.Module) -> None:
        if pretrained_attn is None:
            return

        try:
            for p in pretrained_attn.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        cls_name = pretrained_attn.__class__.__name__
        has_out_lin = any(
            hasattr(pretrained_attn, a)
            for a in ("out_lin", "c_proj", "o_proj", "out_proj", "dense")
        )
        is_nested = False
        if (
            (
                hasattr(pretrained_attn, "self")
                and hasattr(pretrained_attn.self, "query")
                and hasattr(pretrained_attn, "output")
                and hasattr(pretrained_attn.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)

        if hasattr(pretrained_attn, "rotary_emb") and pretrained_attn.rotary_emb is not None:
            self.rotary_emb = pretrained_attn.rotary_emb

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        def _unwrap(m):
            return getattr(m, "base_layer", m)

        # 1. GPT-2 fused c_attn
        if hasattr(pretrained_attn, "c_attn") and pretrained_attn.c_attn is not None:
            c_attn_mod = pretrained_attn.c_attn
            if (
                hasattr(c_attn_mod, "merge")
                and callable(c_attn_mod.merge)
                and not getattr(c_attn_mod, "merged", False)
            ):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    logger.debug("c_attn merge skipped: %s", e)
            base_c = _unwrap(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 hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                )
                if is_conv:
                    c_w = c_w.t().contiguous()
                split = 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], c_b[:split] if c_b is not None else None)
                extracted["k"] = (
                    c_w[split : 2 * split],
                    c_b[split : 2 * split] if c_b is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :],
                    c_b[2 * split :] if c_b is not None else None,
                )

            if hasattr(pretrained_attn, "c_proj") and pretrained_attn.c_proj is not None:
                c_proj_mod = pretrained_attn.c_proj
                if (
                    hasattr(c_proj_mod, "merge")
                    and callable(c_proj_mod.merge)
                    and not getattr(c_proj_mod, "merged", False)
                ):
                    try:
                        c_proj_mod.merge()
                    except Exception as e:
                        logger.debug("c_proj merge skipped: %s", e)
                base_o = _unwrap(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 hasattr(c_proj_mod, "nf")
                        or hasattr(base_o, "nf")
                        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. Nested (BERT / RoBERTa)
        elif hasattr(pretrained_attn, "self") and hasattr(pretrained_attn.self, "query"):
            s = pretrained_attn.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(pretrained_attn, "output") and hasattr(pretrained_attn.output, "dense"):
                d = pretrained_attn.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # 3. Flat modules (DistilBERT, Qwen, LLaMA)
        else:
            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 key, names in candidates.items():
                for name in names:
                    if hasattr(pretrained_attn, name):
                        m = getattr(pretrained_attn, name)
                        if hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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(pretrained_attn, 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:
            logger.warning(
                "EATBlock copy_weights: missing %s; using initialized weights",
                missing_after,
            )

        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

        with torch.no_grad():
            for key, target in (
                ("q", self.q_lin),
                ("k", self.k_lin),
                ("v", self.v_lin),
                ("o", self.out_lin),
            ):
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target.weight)
                        if target.bias is not None:
                            nn.init.zeros_(target.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target.weight.shape
                    and w.t().contiguous().shape == target.weight.shape
                ):
                    w = w.t().contiguous()

                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))
                target.weight.data.zero_()
                target.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                )

                if b is not None and target.bias is not None:
                    mb = min(target.bias.numel(), b.numel())
                    target.bias.data.zero_()
                    target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))
                elif target.bias is not None:
                    nn.init.zeros_(target.bias)

                # EAT-FIX-15: warn (not raise)
                n = float(torch.linalg.norm(target.weight.data.float()).item())
                if n < 1e-4:
                    logger.warning("EAT: weight for '%s' near-zero norm (%.2e)", key, n)

                if self.freeze_base:
                    target.freeze_base_weights()

        # ---- EAT-FIX-09: track compressed ----
        try:
            for m in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

    @classmethod
    def from_distilbert_pretrained(cls, old_attn, model_config=None, **kwargs):
        return cls(config=model_config, pretrained_attn=old_attn, **kwargs)

    # ------------------------------------------------------------------
    # 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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        dtype_in = hidden_states.dtype

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

        q = (
            self.q_lin(hidden_states)
            .view(B, L, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_lin(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_lin(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- RoPE ----
        cos = kwargs.get("cos")
        sin = kwargs.get("sin")
        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) == 2:
            cos, sin = pos_emb[0], pos_emb[1]

        if cos is not None and sin is not None and not is_cross:
            q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        elif self.rotary_emb is not None and not is_cross:
            try:
                pos_ids = position_ids
                if pos_ids is None:
                    pos_ids = (
                        torch.arange(kv_len, device=hidden_states.device).unsqueeze(0).expand(B, -1)
                    )
                cos_, sin_ = self.rotary_emb(k, seq_len=kv_len)
                q, k = apply_rotary_pos_emb(q, k, cos_, sin_, pos_ids)
            except Exception as e:
                logger.debug("EAT: rotary apply failed: %s", e)

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

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        is_causal = bool(self.is_decoder and not is_cross and L > 1)

        # ---- EAT-FIX-05: unified causal + mask ----
        attn_mask_combined = None
        if eff_mask is not None or is_causal:
            if eff_mask is not None:
                if eff_mask.dim() == 2:
                    attn_mask_combined = (eff_mask > 0)[:, None, None, :]  # (B,1,1,N_kv)
                elif eff_mask.dim() == 3:
                    attn_mask_combined = (eff_mask > -1e4)[:, None, :, :]
                elif eff_mask.dim() == 4:
                    attn_mask_combined = eff_mask > -1e4
                else:
                    attn_mask_combined = None

            if is_causal:
                causal = torch.tril(
                    torch.ones(L, kv_len, device=hidden_states.device, dtype=torch.bool)
                )[None, None, :, :]
                if attn_mask_combined is not None:
                    attn_mask_combined = attn_mask_combined & causal
                else:
                    attn_mask_combined = causal

        # ---- Dense attention ----
        scale = 1.0 / math.sqrt(self.head_dim)

        if attn_mask_combined is not None:
            scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
            scores = scores.masked_fill(~attn_mask_combined, -1e4)
            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(dtype_in)
            probs = torch.nan_to_num(probs, nan=0.0)
            probs = self.dropout(probs)
            attn_weights_for_return = probs
        else:
            # Fast SDPA path
            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
                _sdpa_out = F.scaled_dot_product_attention(
                    _qs,
                    _ks,
                    _vs,
                    is_causal=is_causal,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                probs = None
                attn_weights_for_return = None
                _sdpa_out = _sdpa_out.to(dtype_in)
            except Exception:
                scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
                if is_causal:
                    causal = (
                        torch.triu(
                            torch.full(
                                (L, kv_len),
                                -1e4,
                                device=scores.device,
                                dtype=torch.float32,
                            ),
                            diagonal=1 + (kv_len - L),
                        )
                        .unsqueeze(0)
                        .unsqueeze(0)
                    )
                    scores = scores + causal
                scores = torch.nan_to_num(scores, nan=-1e4)
                probs = F.softmax(scores, dim=-1).to(dtype_in)
                probs = torch.nan_to_num(probs, nan=0.0)
                probs = self.dropout(probs)
                attn_weights_for_return = probs
                _sdpa_out = None

        # ---- EAT-FIX-06: head_mask ----
        if head_mask is not None:
            hm = head_mask.to(dtype_in)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)

            if probs is not None:
                probs = probs * hm
            else:
                _sdpa_out = _sdpa_out * hm

        # ---- Context ----
        if probs is not None:
            context = torch.matmul(probs.to(v.dtype), v)
        else:
            context = _sdpa_out

        context = context.transpose(1, 2).contiguous().view(B, L, self.all_head_size)

        # ---- Output projection ----
        # EAT-FIX-03: is_nested_encoder is only True for truly nested containers
        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_lin(context)

        # EAT-FIX-01: no padding/truncation hack
        if out.shape[-1] != D:
            raise RuntimeError(
                f"EATBlock: output dim {out.shape[-1]} != hidden_size {D}. "
                f"Check num_heads * head_dim == hidden_size."
            )

        # ---- Progressive Token Pruning (eval only) ----
        if self.prune_layer is not None and not self.training:
            try:
                out, _ = self.prune_layer(out)
            except Exception as e:
                logger.debug("EAT: pruning skipped: %s", e)

        return out, (attn_weights_for_return if output_attentions else None)

Every 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.

Static methods

def from_distilbert_pretrained(old_attn, model_config=None, **kwargs)

Instance variables

prop c_proj

Expand source code

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

prop dense

Expand source code

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

prop k_proj

Expand source code

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

prop key

Expand source code

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

prop o_proj

Expand source code

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

prop out_proj

Expand source code

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

prop q_proj

Expand source code

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

prop query

Expand source code

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

prop v_proj

Expand source code

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

prop value

Expand source code

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

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class EarlyStoppingCallback (early_stopping_patience: int = 1,
early_stopping_threshold: float | None = 0.0)

Expand source code

class EarlyStoppingCallback(TrainerCallback, ExportableState):
    """
    A [`TrainerCallback`] that handles early stopping.

    Args:
        early_stopping_patience (`int`):
            Use with `metric_for_best_model` to stop training when the specified metric worsens for
            `early_stopping_patience` evaluation calls.
        early_stopping_threshold(`float`, *optional*):
            Use with TrainingArguments `metric_for_best_model` and `early_stopping_patience` to denote how much the
            specified metric must improve to satisfy early stopping conditions. `

    This callback depends on [`TrainingArguments`] argument *load_best_model_at_end* functionality to set best_metric
    in [`TrainerState`]. Note that if the [`TrainingArguments`] argument *save_steps* differs from *eval_steps*, the
    early stopping will not occur until the next save step.
    """

    def __init__(self, early_stopping_patience: int = 1, early_stopping_threshold: float | None = 0.0):
        self.early_stopping_patience = early_stopping_patience
        self.early_stopping_threshold = early_stopping_threshold
        # early_stopping_patience_counter denotes the number of times validation metrics failed to improve.
        self.early_stopping_patience_counter = 0

    def check_metric_value(self, args, state, control, metric_value):
        # best_metric is set by code for load_best_model
        operator = np.greater if args.greater_is_better else np.less
        if state.best_metric is None or (
            operator(metric_value, state.best_metric)
            and abs(metric_value - state.best_metric) > self.early_stopping_threshold
        ):
            self.early_stopping_patience_counter = 0
        else:
            self.early_stopping_patience_counter += 1

    def on_train_begin(self, args, state, control, **kwargs):
        if not args.load_best_model_at_end:
            logger.warning(
                "Using EarlyStoppingCallback without load_best_model_at_end=True. "
                "Once training is finished, the best model will not be loaded automatically."
            )
        assert args.metric_for_best_model is not None, (
            "EarlyStoppingCallback requires metric_for_best_model to be defined"
        )
        assert args.eval_strategy != IntervalStrategy.NO, (
            "EarlyStoppingCallback requires IntervalStrategy of steps or epoch"
        )

    def on_evaluate(self, args, state, control, metrics, **kwargs):
        metric_to_check = args.metric_for_best_model
        if not metric_to_check.startswith("eval_"):
            metric_to_check = f"eval_{metric_to_check}"
        metric_value = metrics.get(metric_to_check)

        if metric_value is None:
            logger.warning(
                f"early stopping required metric_for_best_model, but did not find {metric_to_check} so early stopping"
                " is disabled"
            )
            return

        self.check_metric_value(args, state, control, metric_value)
        if self.early_stopping_patience_counter >= self.early_stopping_patience:
            control.should_training_stop = True

    def state(self) -> dict:
        return {
            "args": {
                "early_stopping_patience": self.early_stopping_patience,
                "early_stopping_threshold": self.early_stopping_threshold,
            },
            "attributes": {
                "early_stopping_patience_counter": self.early_stopping_patience_counter,
            },
        }

A [TrainerCallback] that handles early stopping.

Args -----= early_stopping_patience (int): Use with metric_for_best_model to stop training when the specified metric worsens for early_stopping_patience evaluation calls. early_stopping_threshold(float, optional): Use with TrainingArguments metric_for_best_model and early_stopping_patience to denote how much the specified metric must improve to satisfy early stopping conditions. This callback depends on [TrainingArguments] argument *load_best_model_at_end* functionality to set best_metric in [TrainerState]. Note that if the [TrainingArguments`] argument save_steps differs from eval_steps, the early stopping will not occur until the next save step.

Ancestors

• transformers.trainer_callback.TrainerCallback
• transformers.trainer_callback.ExportableState

Methods

def check_metric_value(self, args, state, control, metric_value)

Expand source code

def check_metric_value(self, args, state, control, metric_value):
    # best_metric is set by code for load_best_model
    operator = np.greater if args.greater_is_better else np.less
    if state.best_metric is None or (
        operator(metric_value, state.best_metric)
        and abs(metric_value - state.best_metric) > self.early_stopping_threshold
    ):
        self.early_stopping_patience_counter = 0
    else:
        self.early_stopping_patience_counter += 1

def on_evaluate(self, args, state, control, metrics, **kwargs)

Expand source code

def on_evaluate(self, args, state, control, metrics, **kwargs):
    metric_to_check = args.metric_for_best_model
    if not metric_to_check.startswith("eval_"):
        metric_to_check = f"eval_{metric_to_check}"
    metric_value = metrics.get(metric_to_check)

    if metric_value is None:
        logger.warning(
            f"early stopping required metric_for_best_model, but did not find {metric_to_check} so early stopping"
            " is disabled"
        )
        return

    self.check_metric_value(args, state, control, metric_value)
    if self.early_stopping_patience_counter >= self.early_stopping_patience:
        control.should_training_stop = True

Event called after an evaluation phase.

def on_train_begin(self, args, state, control, **kwargs)

Expand source code

def on_train_begin(self, args, state, control, **kwargs):
    if not args.load_best_model_at_end:
        logger.warning(
            "Using EarlyStoppingCallback without load_best_model_at_end=True. "
            "Once training is finished, the best model will not be loaded automatically."
        )
    assert args.metric_for_best_model is not None, (
        "EarlyStoppingCallback requires metric_for_best_model to be defined"
    )
    assert args.eval_strategy != IntervalStrategy.NO, (
        "EarlyStoppingCallback requires IntervalStrategy of steps or epoch"
    )

Event called at the beginning of training.

def state(self) ‑> dict

Expand source code

def state(self) -> dict:
    return {
        "args": {
            "early_stopping_patience": self.early_stopping_patience,
            "early_stopping_threshold": self.early_stopping_threshold,
        },
        "attributes": {
            "early_stopping_patience_counter": self.early_stopping_patience_counter,
        },
    }

class ExpertChoiceMoE (d_model: int,
n_experts: int = 2,
capacity_factor: float = 1.0,
noise_std: float = 0.05)

Expand source code

class ExpertChoiceMoE(nn.Module):
    r"""
    Expert Choice Routing MoE layer.
    Each expert selects top-k tokens with the highest affinity scores.
    Equipped with zero-initialized final layers for step-0 representation preservation.
    """

    def __init__(
        self,
        d_model: int,
        n_experts: int = 2,
        capacity_factor: float = 1.0,
        noise_std: float = 0.05,
    ) -> None:
        super().__init__()
        self.d_model = d_model
        self.n_experts = n_experts
        self.capacity_factor = capacity_factor
        self.noise_std = noise_std
        self.aux_loss: torch.Tensor = torch.tensor(0.0)

        self.router_ln = nn.LayerNorm(self.d_model)
        self.router = nn.Linear(d_model, n_experts, bias=False)
        nn.init.uniform_(self.router.weight, -0.01, 0.01)

        expert_hidden_dim = max(32, d_model // 4)
        self.experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(d_model, expert_hidden_dim),
                    nn.GELU(),
                    nn.Linear(expert_hidden_dim, d_model),
                )
                for _ in range(n_experts)
            ]
        )

        # Active MoE expert representation
        for expert in self.experts:
            nn.init.normal_(expert[2].weight, std=0.02)
            if expert[2].bias is not None:
                nn.init.zeros_(expert[2].bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        r"""Forward execution using Expert Choice routing."""
        B, L, D = x.shape
        x_flat = x.view(-1, D)
        T = x_flat.size(0)

        if self.d_model == D:
            logits = self.router(self.router_ln(x_flat))
        else:
            normed = F.layer_norm(x_flat, (D,))
            if not hasattr(self, "_dynamic_router") or self._dynamic_router.in_features != D:
                self._dynamic_router = nn.Linear(D, self.n_experts, bias=False).to(
                    device=x.device, dtype=x.dtype
                )
                nn.init.uniform_(self._dynamic_router.weight, -0.01, 0.01)
            logits = self._dynamic_router(normed)

        if self.training and self.noise_std > 0:
            logits = logits + torch.randn_like(logits) * self.noise_std

        # Softmax over tokens per expert
        gate_scores = F.softmax(logits.float(), dim=0).to(dtype=logits.dtype)
        k = max(1, min(T, int(T * self.capacity_factor / self.n_experts)))

        out_flat = torch.zeros_like(x_flat)
        for e in range(self.n_experts):
            scores_e = gate_scores[:, e]
            topk_vals, topk_indices = torch.topk(scores_e, k, dim=0)

            selected_x = x_flat[topk_indices]
            if selected_x.size(-1) == self.d_model:
                expert_out = self.experts[e](selected_x)
            else:
                if (
                    not hasattr(self, f"_expert_in_proj_{e}")
                    or getattr(self, f"_expert_in_proj_{e}").in_features != D
                ):
                    setattr(
                        self,
                        f"_expert_in_proj_{e}",
                        nn.Linear(D, self.d_model).to(device=x.device, dtype=x.dtype),
                    )
                    setattr(
                        self,
                        f"_expert_out_proj_{e}",
                        nn.Linear(self.d_model, D).to(device=x.device, dtype=x.dtype),
                    )
                proj_in = getattr(self, f"_expert_in_proj_{e}")
                proj_out = getattr(self, f"_expert_out_proj_{e}")
                expert_out = proj_out(self.experts[e](proj_in(selected_x)))

            update_val = (
                expert_out.to(out_flat.dtype) * topk_vals.unsqueeze(-1).to(out_flat.dtype)
            ).to(dtype=out_flat.dtype)
            out_flat.index_add_(0, topk_indices, update_val)

        # Auxiliary load balancing loss
        prob_mean = F.softmax(logits.float() / 1.0, dim=-1).mean(dim=0)
        self.aux_loss = self.n_experts * torch.sum(prob_mean**2)

        return out_flat.view(B, L, D)

Expert Choice Routing MoE layer. Each expert selects top-k tokens with the highest affinity scores. Equipped with zero-initialized final layers for step-0 representation preservation.

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:
    r"""Forward execution using Expert Choice routing."""
    B, L, D = x.shape
    x_flat = x.view(-1, D)
    T = x_flat.size(0)

    if self.d_model == D:
        logits = self.router(self.router_ln(x_flat))
    else:
        normed = F.layer_norm(x_flat, (D,))
        if not hasattr(self, "_dynamic_router") or self._dynamic_router.in_features != D:
            self._dynamic_router = nn.Linear(D, self.n_experts, bias=False).to(
                device=x.device, dtype=x.dtype
            )
            nn.init.uniform_(self._dynamic_router.weight, -0.01, 0.01)
        logits = self._dynamic_router(normed)

    if self.training and self.noise_std > 0:
        logits = logits + torch.randn_like(logits) * self.noise_std

    # Softmax over tokens per expert
    gate_scores = F.softmax(logits.float(), dim=0).to(dtype=logits.dtype)
    k = max(1, min(T, int(T * self.capacity_factor / self.n_experts)))

    out_flat = torch.zeros_like(x_flat)
    for e in range(self.n_experts):
        scores_e = gate_scores[:, e]
        topk_vals, topk_indices = torch.topk(scores_e, k, dim=0)

        selected_x = x_flat[topk_indices]
        if selected_x.size(-1) == self.d_model:
            expert_out = self.experts[e](selected_x)
        else:
            if (
                not hasattr(self, f"_expert_in_proj_{e}")
                or getattr(self, f"_expert_in_proj_{e}").in_features != D
            ):
                setattr(
                    self,
                    f"_expert_in_proj_{e}",
                    nn.Linear(D, self.d_model).to(device=x.device, dtype=x.dtype),
                )
                setattr(
                    self,
                    f"_expert_out_proj_{e}",
                    nn.Linear(self.d_model, D).to(device=x.device, dtype=x.dtype),
                )
            proj_in = getattr(self, f"_expert_in_proj_{e}")
            proj_out = getattr(self, f"_expert_out_proj_{e}")
            expert_out = proj_out(self.experts[e](proj_in(selected_x)))

        update_val = (
            expert_out.to(out_flat.dtype) * topk_vals.unsqueeze(-1).to(out_flat.dtype)
        ).to(dtype=out_flat.dtype)
        out_flat.index_add_(0, topk_indices, update_val)

    # Auxiliary load balancing loss
    prob_mean = F.softmax(logits.float() / 1.0, dim=-1).mean(dim=0)
    self.aux_loss = self.n_experts * torch.sum(prob_mean**2)

    return out_flat.view(B, L, D)

Forward execution using Expert Choice routing.

class ExpertChoiceMoEAttention (config: Any | None = None,
rank: int | None = 64,
n_experts: int = 2,
capacity_factor: float = 1.0,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
moe_init_gate: float = 0.0,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "expert_choice_moe",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "expert_choice_moe_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class ExpertChoiceMoEAttention(BaseAttention):
    r"""Expert Choice MoE Multi-Head Attention Module."""

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        n_experts: int = 2,
        capacity_factor: float = 1.0,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        moe_init_gate: float = 0.0,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        cfg_obj = getattr(config, "text_config", config) if config is not None else None
        hidden_size = getattr(cfg_obj, "hidden_size", None) if cfg_obj is not None else None
        num_heads = getattr(cfg_obj, "num_attention_heads", None) if cfg_obj is not None else None

        if hidden_size is None and source_attn is not None:
            for cand in ["q_lin", "query", "q_proj", "c_attn", "in_proj_weight", "qkv"]:
                mod = getattr(source_attn, cand, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    hidden_size = w.shape[1] if w.ndim == 2 else None
                    break
                elif isinstance(mod, torch.Tensor) and mod.ndim == 2:
                    hidden_size = mod.shape[1]
                    break
        if hidden_size is None:
            hidden_size = kwargs.get("hidden_size", 768)
        if num_heads is None:
            num_heads = kwargs.get("num_attention_heads", 12)

        super().__init__(config, rank)
        self.dim = int(hidden_size)
        self.num_heads = int(num_heads)
        self.num_kv_heads = (
            getattr(config, "num_key_value_heads", self.num_heads)
            if config is not None
            else self.num_heads
        )
        self.head_dim = getattr(config, "head_dim", None) or (self.dim // max(1, self.num_heads))
        self.all_head_size = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))

        # Static layer index
        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx: int = int(resolved_layer_id)

        # LoRA adaptation configuration
        self.lora_rank: int = getattr(
            config,
            "lora_rank",
            getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha: float = getattr(
            config,
            "lora_alpha",
            getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base: bool = self.lora_rank > 0

        # Decoder & nested encoder detection
        model_type = getattr(config, "model_type", "").lower()
        is_dec = getattr(
            config,
            "is_decoder",
            getattr(config, "is_causal", kwargs.get("is_causal", False)),
        )
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(model_type.startswith(p) for p in decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)
        self.is_nested_encoder: bool = False

        # Multi-head projections equipped with LoRA adapters
        self.q_lin = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_lin = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_lin = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_lin = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.moe = ExpertChoiceMoE(
            self.all_head_size, n_experts=n_experts, capacity_factor=capacity_factor
        )
        # Trainable gate for MoE output contribution (initialized with negative bias for smooth ramp-up)
        self.moe_gate = nn.Parameter(torch.tensor(float(moe_init_gate)))

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)
        self.aux_loss: torch.Tensor = torch.tensor(0.0)

        self.rotary_emb = None
        self.comp_config: Any | None = comp_config

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

    # Compatibility properties
    @property
    def query(self) -> LoRALinear:
        return self.q_lin

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

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

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

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

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

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

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

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

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

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

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

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

    def copy_weights(self, pretrained_attn: nn.Module) -> None:
        r"""Copy projection weights from pretrained attention module with Frobenius norm verification."""
        if pretrained_attn is None:
            return

        cls_name = pretrained_attn.__class__.__name__
        has_out_lin = (
            hasattr(pretrained_attn, "out_lin")
            or hasattr(pretrained_attn, "c_proj")
            or hasattr(pretrained_attn, "o_proj")
        )
        is_nested = False
        if (
            (
                hasattr(pretrained_attn, "self")
                and hasattr(pretrained_attn.self, "query")
                and hasattr(pretrained_attn, "output")
                and hasattr(pretrained_attn.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)

        if hasattr(pretrained_attn, "rotary_emb"):
            self.rotary_emb = pretrained_attn.rotary_emb

        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 + c_proj
        if hasattr(pretrained_attn, "c_attn") and pretrained_attn.c_attn is not None:
            c_attn_mod = pretrained_attn.c_attn
            if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    print(f"copy_weights failed: {e}")
                    logger.warning(f"copy_weights failed: {e}")
            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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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(pretrained_attn, "c_proj") and pretrained_attn.c_proj is not None:
                c_proj_mod = pretrained_attn.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()
                    if (
                        (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])
                    ):
                        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. RoBERTa / BERT nested self-attention
        elif hasattr(pretrained_attn, "self") and hasattr(pretrained_attn.self, "query"):
            s = pretrained_attn.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(pretrained_attn, "output") and hasattr(pretrained_attn.output, "dense"):
                out_mod = pretrained_attn.output.dense
                extracted["o"] = (
                    out_mod.weight.data.clone(),
                    out_mod.bias.data.clone() if out_mod.bias is not None else None,
                )

        # 3. Standard linear projections (DistilBERT, Qwen, LLaMA)
        else:
            proj_candidates = {
                "q": ["q_lin", "q_proj", "query", "W_q"],
                "k": ["k_lin", "k_proj", "key", "W_k"],
                "v": ["v_lin", "v_proj", "value", "W_v"],
                "o": ["out_lin", "out_proj", "dense", "o_proj"],
            }
            for proj_key, names in proj_candidates.items():
                for name in names:
                    if hasattr(pretrained_attn, name):
                        mod = getattr(pretrained_attn, 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

        # Fallback via UniversalAdapter
        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(pretrained_attn, 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:
            logger.warning(
                f"ExpertChoiceMoEAttention copy_weights failed: missing mandatory projections {missing_after} "
                f"from module {pretrained_attn.__class__.__name__}. Proceeding with partially initialized projections."
            )

        target_map = {
            "q": self.q_lin,
            "k": self.k_lin,
            "v": self.v_lin,
            "o": self.out_lin,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

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

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape

        is_cross_attention = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross_attention else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Projections
        q = (
            self.q_lin(hidden_states)
            .view(B, L, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_lin(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_lin(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. Rotary Position Embeddings (RoPE)
        if position_embeddings is not None and not is_cross_attention:
            if isinstance(position_embeddings, tuple):
                cos, sin = position_embeddings
            else:
                cos, sin = position_embeddings, None
            if sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        elif self.rotary_emb is not None and not is_cross_attention:
            if position_ids is not None:
                cos, sin = self.rotary_emb(v, seq_len=kv_len)
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)

        # 3. Grouped Query Attention (GQA) head expansion
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        # 4. Multi-head attention scores
        scale = 1.0 / math.sqrt(self.head_dim)
        # Memory-efficient: use SDPA to avoid OOM on large seqs

        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 = getattr(self, "is_decoder", False) and not is_cross_attention and L > 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,
            )
            attn_out = (
                _sdpa_out.to(hidden_states.dtype)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.all_head_size)
            )
            attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
            moe_out = self.moe(attn_out)
            self.aux_loss = self.moe.aux_loss
            moe_weight = torch.sigmoid(self.moe_gate).to(attn_out.dtype)
            enriched_attn = (attn_out + moe_weight * moe_out.to(attn_out.dtype)).to(attn_out.dtype)
            if self.is_nested_encoder and enriched_attn.shape[-1] == self.dim:
                out = enriched_attn
            else:
                out = self.out_lin(enriched_attn)
            if out.shape[-1] != hidden_states.shape[-1]:
                out = (
                    F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
                    if out.shape[-1] < hidden_states.shape[-1]
                    else out[..., : hidden_states.shape[-1]]
                )
            return out, None
        except Exception:
            pass

        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross_attention:
            causal_mask = torch.triu(
                torch.full(
                    (L, kv_len),
                    float("-inf"),
                    device=scores.device,
                    dtype=torch.float32,
                ),
                diagonal=1,
            )
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

        eff_mask = encoder_attention_mask if is_cross_attention else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                pad_mask = (eff_mask == 0).unsqueeze(1).unsqueeze(2)
                scores = scores.masked_fill(pad_mask, float("-inf"))
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                scores = scores + eff_mask.float()

        probs = F.softmax(scores, dim=-1)
        probs = torch.nan_to_num(probs, nan=0.0)
        probs = self.dropout(probs)

        # 5. Head mask handling
        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.shape[0], head_mask.shape[1], 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            probs = probs * head_mask

        attn_out = (
            torch.matmul(probs.to(v.dtype), v)
            .transpose(1, 2)
            .contiguous()
            .view(B, L, self.all_head_size)
        )

        # 6. MoE feedforward enrichment with gate blending
        moe_out = self.moe(attn_out)
        self.aux_loss = self.moe.aux_loss
        moe_weight = torch.sigmoid(self.moe_gate).to(attn_out.dtype)
        enriched_attn = (attn_out + moe_weight * moe_out.to(attn_out.dtype)).to(attn_out.dtype)

        # 7. Output projection with nested encoder bypass
        if self.is_nested_encoder:
            if enriched_attn.shape[-1] != self.dim:
                out = self.out_lin(enriched_attn)
            else:
                out = enriched_attn
        else:
            out = self.out_lin(enriched_attn)

        if out.shape[-1] != hidden_states.shape[-1]:
            out = (
                F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
                if out.shape[-1] < hidden_states.shape[-1]
                else out[..., : hidden_states.shape[-1]]
            )

        return out, (probs if output_attentions else None)

Expert Choice MoE Multi-Head Attention Module.

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.

Instance variables

prop W_k : LoRALinear

Expand source code

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

prop W_q : LoRALinear

Expand source code

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

prop W_v : LoRALinear

Expand source code

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

prop c_proj : LoRALinear

Expand source code

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

prop dense : LoRALinear

Expand source code

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

prop k_proj : LoRALinear

Expand source code

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

prop key : LoRALinear

Expand source code

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

prop o_proj : LoRALinear

Expand source code

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

prop out_proj : LoRALinear

Expand source code

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

prop q_proj : LoRALinear

Expand source code

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

prop query : LoRALinear

Expand source code

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

prop v_proj : LoRALinear

Expand source code

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

prop value : LoRALinear

Expand source code

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

Methods

def copy_weights(self, pretrained_attn: nn.Module) ‑> None

Expand source code

def copy_weights(self, pretrained_attn: nn.Module) -> None:
    r"""Copy projection weights from pretrained attention module with Frobenius norm verification."""
    if pretrained_attn is None:
        return

    cls_name = pretrained_attn.__class__.__name__
    has_out_lin = (
        hasattr(pretrained_attn, "out_lin")
        or hasattr(pretrained_attn, "c_proj")
        or hasattr(pretrained_attn, "o_proj")
    )
    is_nested = False
    if (
        (
            hasattr(pretrained_attn, "self")
            and hasattr(pretrained_attn.self, "query")
            and hasattr(pretrained_attn, "output")
            and hasattr(pretrained_attn.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)

    if hasattr(pretrained_attn, "rotary_emb"):
        self.rotary_emb = pretrained_attn.rotary_emb

    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 + c_proj
    if hasattr(pretrained_attn, "c_attn") and pretrained_attn.c_attn is not None:
        c_attn_mod = pretrained_attn.c_attn
        if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
            try:
                c_attn_mod.merge()
            except Exception as e:
                print(f"copy_weights failed: {e}")
                logger.warning(f"copy_weights failed: {e}")
        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()
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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(pretrained_attn, "c_proj") and pretrained_attn.c_proj is not None:
            c_proj_mod = pretrained_attn.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()
                if (
                    (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])
                ):
                    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. RoBERTa / BERT nested self-attention
    elif hasattr(pretrained_attn, "self") and hasattr(pretrained_attn.self, "query"):
        s = pretrained_attn.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(pretrained_attn, "output") and hasattr(pretrained_attn.output, "dense"):
            out_mod = pretrained_attn.output.dense
            extracted["o"] = (
                out_mod.weight.data.clone(),
                out_mod.bias.data.clone() if out_mod.bias is not None else None,
            )

    # 3. Standard linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q": ["q_lin", "q_proj", "query", "W_q"],
            "k": ["k_lin", "k_proj", "key", "W_k"],
            "v": ["v_lin", "v_proj", "value", "W_v"],
            "o": ["out_lin", "out_proj", "dense", "o_proj"],
        }
        for proj_key, names in proj_candidates.items():
            for name in names:
                if hasattr(pretrained_attn, name):
                    mod = getattr(pretrained_attn, 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

    # Fallback via UniversalAdapter
    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(pretrained_attn, 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:
        logger.warning(
            f"ExpertChoiceMoEAttention copy_weights failed: missing mandatory projections {missing_after} "
            f"from module {pretrained_attn.__class__.__name__}. Proceeding with partially initialized projections."
        )

    target_map = {
        "q": self.q_lin,
        "k": self.k_lin,
        "v": self.v_lin,
        "o": self.out_lin,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device, dtype=target_mod.weight.dtype
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

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

            if self.freeze_base:
                target_mod.freeze_base_weights()

Copy projection weights from pretrained attention module with Frobenius norm verification.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class FUSIONATTN (config: Any | None = None,
hidden_size: int = 768,
num_heads: int = 12,
block_size: int = 64,
top_k_blocks: int = 4,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "fusion_attn", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "fusion_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class FUSIONATTN(BaseAttention):
    """
    Hybrid sparse attention with block gating.

    Two execution paths:
      * Short sequence (N <= 128) or cross-attention  → dense attention
      * Long sequence                                  → block-gated sparse
    """

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

    def __init__(
        self,
        config: Any | None = None,
        hidden_size: int = 768,
        num_heads: int = 12,
        block_size: int = 64,
        top_k_blocks: int = 4,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # ---- resolve hidden_size / num_heads ----
        if config is not None:
            hidden_size = int(
                getattr(
                    config,
                    "hidden_size",
                    getattr(config, "d_model", getattr(config, "n_embd", hidden_size)),
                )
            )
            num_heads = int(
                getattr(
                    config,
                    "num_attention_heads",
                    getattr(config, "num_heads", getattr(config, "n_heads", num_heads)),
                )
            )
        elif source_attn is not None:
            out_layer = None
            for attr in ("out_lin", "dense", "c_proj", "o_proj", "out_proj"):
                if hasattr(source_attn, attr):
                    out_layer = getattr(source_attn, attr)
                    break
            if out_layer is not None:
                hidden_size = int(
                    getattr(
                        out_layer,
                        "out_features",
                        getattr(out_layer, "in_features", hidden_size),
                    )
                )
            num_heads = int(
                getattr(
                    source_attn,
                    "n_heads",
                    getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(source_attn, "num_heads", num_heads),
                    ),
                )
            )

        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        self.hidden_size = int(hidden_size)
        self.num_heads = max(1, int(num_heads))

        # FA-FIX-13: num_kv_heads fallback chain
        if config is not None:
            self.num_kv_heads = int(
                getattr(
                    config,
                    "num_key_value_heads",
                    getattr(config, "num_kv_heads", self.num_heads),
                )
            )
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))

        cfg_hd = getattr(config, "head_dim", None) if config else None
        self.head_dim = (
            int(cfg_hd) if cfg_hd is not None else max(1, self.hidden_size // 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 // self.num_kv_heads)

        self.block_size = int(block_size)
        self.top_k_blocks = int(top_k_blocks)

        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx = int(resolved_layer_id)

        # ---- LoRA ----
        self.lora_rank = int(
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else kwargs.get("lora_rank", 16)
        )
        self.lora_alpha = float(
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else kwargs.get("lora_alpha", 32.0)
        )
        self.freeze_base = self.lora_rank > 0

        # ---- decoder detection ----
        model_type = str(getattr(config, "model_type", "") if config else "").lower()
        is_dec = False
        if config is not None:
            is_dec = bool(
                getattr(config, "is_decoder", False) or getattr(config, "is_causal", False)
            )
        _dec_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if any(model_type.startswith(p) for p in _dec_prefixes):
            is_dec = True
        if config is not None and hasattr(config, "architectures"):
            arch = getattr(config, "architectures", []) or []
            if any("GPT" in str(a) or "CausalLM" in str(a) for a in arch):
                is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        if source_attn is not None:
            if hasattr(source_attn, "c_attn") or getattr(source_attn, "is_decoder", False):
                is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # ---- projections ----
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )

        # FA-FIX-15: per-head gate MLP (head_id embedding + centroid features)
        self.attn_gate = nn.Sequential(
            nn.Linear(self.head_dim * 2, 4 * self.head_dim),
            nn.GELU(),
            nn.Linear(4 * self.head_dim, 1),
        )
        # learnable per-head bias so heads can diverge
        self.gate_head_bias = nn.Parameter(torch.zeros(self.num_heads))
        self.gate_temp = nn.Parameter(torch.ones(1))

        dropout_p = (
            float(
                getattr(
                    config,
                    "dropout",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                )
            )
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )
        self.comp_config = comp_config

        if source_attn is not None:
            self.copy_weights(source_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 o_proj(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 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 c_proj(self):
        return self.out_proj

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, pretrained_attn: nn.Module) -> None:
        if pretrained_attn is None:
            return

        cls_name = pretrained_attn.__class__.__name__
        has_out_lin = any(hasattr(pretrained_attn, a) for a in ("out_lin", "c_proj", "o_proj"))
        is_nested = False
        if (
            (
                hasattr(pretrained_attn, "self")
                and hasattr(pretrained_attn.self, "query")
                and hasattr(pretrained_attn, "output")
                and hasattr(pretrained_attn.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)

        if hasattr(pretrained_attn, "rotary_emb"):
            self.rotary_emb = pretrained_attn.rotary_emb

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused c_attn
        if hasattr(pretrained_attn, "c_attn") and pretrained_attn.c_attn is not None:
            c_attn_mod = pretrained_attn.c_attn
            if (
                hasattr(c_attn_mod, "merge")
                and callable(c_attn_mod.merge)
                and not getattr(c_attn_mod, "merged", False)
            ):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    logger.debug("c_attn merge skipped: %s", e)

            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()
                cls_c = c_attn_mod.__class__.__name__
                base_cls_c = base_c.__class__.__name__
                is_conv = (
                    "Conv1D" in cls_c
                    or "Conv1D" in base_cls_c
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                )
                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(pretrained_attn, "c_proj") and pretrained_attn.c_proj is not None:
                c_proj_mod = pretrained_attn.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()
                    cls_o = c_proj_mod.__class__.__name__
                    base_cls_o = base_o.__class__.__name__
                    is_conv_o = (
                        "Conv1D" in cls_o
                        or "Conv1D" in base_cls_o
                        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,
                    )

        # RoBERTa / BERT nested
        elif hasattr(pretrained_attn, "self") and hasattr(pretrained_attn.self, "query"):
            s = pretrained_attn.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                mod = getattr(s, sub, None)
                if mod is not None and hasattr(mod, "weight"):
                    b = getattr(mod, "bias", None)
                    extracted[key] = (
                        mod.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(pretrained_attn, "output") and hasattr(pretrained_attn.output, "dense"):
                out_mod = pretrained_attn.output.dense
                b = getattr(out_mod, "bias", None)
                extracted["o"] = (
                    out_mod.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat modules
        else:
            candidates = {
                "q": ("q_lin", "q_proj", "query", "W_q"),
                "k": ("k_lin", "k_proj", "key", "W_k"),
                "v": ("v_lin", "v_proj", "value", "W_v"),
                "o": ("out_lin", "out_proj", "dense", "o_proj"),
            }
            for key, names in candidates.items():
                for n in names:
                    if hasattr(pretrained_attn, n):
                        mod = getattr(pretrained_attn, n)
                        if mod is not None and hasattr(mod, "weight"):
                            b = getattr(mod, "bias", None)
                            extracted[key] = (
                                mod.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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(pretrained_attn, 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:
            logger.warning("FUSIONATTN copy_weights: missing %s from %s", missing_after, cls_name)

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()

                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))
                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target_mod.weight.device, target_mod.weight.dtype)
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(target_mod.bias.device, target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                w_norm = float(torch.norm(target_mod.weight.data).item())
                if w_norm < 1e-4:
                    logger.warning("FUSIONATTN: '%s' near-zero norm (%.2e)", proj_key, w_norm)

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    # ------------------------------------------------------------------
    # KV-cache helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _read_past_kv(past_key_value: Any | None, layer_idx: int = 0):
        if past_key_value is None:
            return None, None
        if hasattr(past_key_value, "key_cache") and hasattr(past_key_value, "value_cache"):
            keys = past_key_value.key_cache
            vals = past_key_value.value_cache
            if len(keys) == 0:
                return None, None
            idx = min(layer_idx, len(keys) - 1)
            k, v = keys[idx], vals[idx]
            if isinstance(k, torch.Tensor) and isinstance(v, torch.Tensor):
                return k, v
        if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
            k, v = past_key_value[0], past_key_value[1]
            if isinstance(k, torch.Tensor) and isinstance(v, torch.Tensor):
                return k, v
        return None, None

    # ------------------------------------------------------------------
    # Mask preparation
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_valid_mask(mask, B: int, N: int, device: torch.device) -> torch.Tensor | None:
        """Return (B, N) bool mask where True = valid token."""
        if mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if mask.dim() == 2:
            m = _to_bool(mask)
            if m.shape[-1] == N:
                return m
        elif mask.dim() == 4:
            m = _to_bool(mask)
            m = m.any(dim=-2).squeeze(1)
            if m.shape[-1] == N:
                return m
        return None

    # ------------------------------------------------------------------
    # Dense path
    # ------------------------------------------------------------------
    def _dense_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        valid_mask: torch.Tensor | None,
        head_mask: torch.Tensor | None,
        is_causal: bool,
        output_attentions: bool,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, H, N, D = q.shape
        N_kv = k.shape[2]
        scale = D**-0.5

        # FA-FIX-02: only use SDPA when no extra features are needed
        use_sdpa = (
            not output_attentions
            and head_mask is None
            and (valid_mask is None or valid_mask.all())
            and not self.training  # training dropout in SDPA is fine, but keep manual for safety
        )

        if use_sdpa:
            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
                _out = F.scaled_dot_product_attention(
                    _qs,
                    _ks,
                    _vs,
                    is_causal=is_causal and N > 1 and N_kv == N,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                _out = _out.to(q.dtype).transpose(1, 2).contiguous().view(B, N, self.all_head_size)
                _out = torch.nan_to_num(_out, nan=0.0, posinf=1e4, neginf=-1e4)
                return _out, None
            except Exception as e:
                logger.debug("SDPA failed, falling back to manual: %s", e)

        # ---- Manual path ----
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        # FA-FIX-03 / decoder: causal mask with past_len offset
        if is_causal and N > 1:
            past_len = max(0, N_kv - N)
            causal = (
                torch.triu(
                    torch.full((N, N_kv), float("-inf"), device=q.device, dtype=torch.float32),
                    diagonal=1 + past_len,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(v.dtype)
        probs = torch.nan_to_num(probs, nan=0.0)
        probs = self.dropout(probs)

        if head_mask is not None:
            hm = head_mask.to(probs.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            probs = probs * hm

        context = torch.matmul(probs, v).transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        return context, (probs if output_attentions else None)

    # ------------------------------------------------------------------
    # Block-gated sparse attention
    # ------------------------------------------------------------------
    def _block_sparse_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        valid_mask: torch.Tensor | None,
        head_mask: torch.Tensor | None,
        is_causal: bool,
        output_attentions: bool,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        FA-FIX-03 / FA-FIX-04 / FA-FIX-05 / FA-FIX-15 / FA-FIX-16 /
        FA-FIX-17 / FA-FIX-18 rewritten block-gated attention.
        """
        B, H, N, D = q.shape
        k.shape[2]
        bs = self.block_size
        device = q.device

        # Pad to multiple of block_size
        pad_len = (bs - (N % bs)) % bs
        if pad_len > 0:
            q_pad = F.pad(q, (0, 0, 0, pad_len))
            k_pad = F.pad(k, (0, 0, 0, pad_len))
            v_pad = F.pad(v, (0, 0, 0, pad_len))
            if valid_mask is not None:
                valid_pad = F.pad(valid_mask, (0, pad_len), value=False)
            else:
                valid_pad = None
        else:
            q_pad, k_pad, v_pad = q, k, v
            valid_pad = valid_mask

        N_pad = q_pad.size(2)
        num_blocks = N_pad // bs
        top_k = min(self.top_k_blocks, num_blocks)

        # ---- block centroids ----
        q_blocks = q_pad.reshape(B, H, num_blocks, bs, D)
        k_blocks = k_pad.reshape(B, H, num_blocks, bs, D)
        q_centroid = q_blocks.mean(dim=3)  # (B, H, nb, D)
        k_centroid = k_blocks.mean(dim=3)  # (B, H, nb, D)

        # FA-FIX-15 / FA-FIX-16: per-head gating with centroid features
        q_feat = torch.cat([q_centroid, torch.zeros_like(q_centroid)], dim=-1)  # (B,H,nb,2D)
        k_feat = torch.cat([torch.zeros_like(k_centroid), k_centroid], dim=-1)  # (B,H,nb,2D)
        q_gate = self.attn_gate(q_feat).squeeze(-1)  # (B, H, nb)
        k_gate = self.attn_gate(k_feat).squeeze(-1)  # (B, H, nb)

        # per-head learnable bias
        head_bias = self.gate_head_bias.view(1, H, 1)
        gate_scores = (
            torch.matmul(q_gate.unsqueeze(-1), k_gate.unsqueeze(-2)).squeeze(-1).squeeze(-1)
            + head_bias
        ) / (self.gate_temp.abs() + 1e-5)
        # gate_scores: (B, H, nb)

        # FA-FIX-18: causal-aware top-k for decoder self-attention
        # A query in block i cannot attend to any key in block j > i.
        if is_causal:
            blk_idx = torch.arange(num_blocks, device=device)
            causal_blocks = blk_idx.unsqueeze(0) <= blk_idx.unsqueeze(1)  # (nb, nb)
            # mask future blocks with -inf so top-k never picks them
            gate_scores.unsqueeze(-1).expand(-1, -1, -1, num_blocks)  # not needed; use simpler form
            # compute per-query-block mask and apply to logits
            # gate_scores is (B, H, nb). We need (B, H, nb) * causal_mask_per_qblock
            # A more direct approach: for each query block i, only keep scores over j <= i
            # which is a triangular mask per (B,H,i).
            # We expand gate_scores to (B, H, nb, nb) with broadcast, then mask.
            gs2 = gate_scores.unsqueeze(-1).expand(
                -1, -1, -1, num_blocks
            )  # (B,H,nb,nb) where last dim = key block
            causal_blocks_b = causal_blocks.view(1, 1, num_blocks, num_blocks)
            gs2 = gs2.masked_fill(~causal_blocks_b, float("-inf"))
            # Top-k over key-block dim
            _, top_indices = torch.topk(gs2, top_k, dim=-1)  # (B, H, nb, top_k)
        else:
            # bidir: broadcast scores over query-block dim
            gs2 = gate_scores.unsqueeze(-1).expand(-1, -1, num_blocks, num_blocks)
            _, top_indices = torch.topk(gs2, top_k, dim=-1)

        # FA-FIX-05: guarantee block 0 is kept by constructing top-k differently
        # For each query block, force slot 0 to be block 0 (sink), fill rest with top-(k-1)
        with torch.no_grad():
            if top_k > 1:
                # take top-(k-1) from scores excluding block 0
                gs2_others = gs2.clone()
                gs2_others[..., 0] = float("-inf")
                _, top_others = torch.topk(gs2_others, top_k - 1, dim=-1)
                sink = torch.zeros_like(top_others[..., :1])  # block 0
                top_indices = torch.cat([sink, top_others], dim=-1)
            else:
                top_indices = torch.zeros(B, H, num_blocks, top_k, device=device, dtype=torch.long)

        # ---- gather K/V per query block ----
        # k_blocks: (B, H, nb, bs, D), v_blocks same
        v_blocks = v_pad.reshape(B, H, num_blocks, bs, D)

        # For padding mask of K, pad to full length then reshape to blocks
        if valid_pad is not None:
            valid_blocks = valid_pad.view(B, 1, num_blocks, bs)  # (B, 1, nb, bs)
        else:
            valid_blocks = torch.ones(B, 1, num_blocks, bs, device=device, dtype=torch.bool)

        # gather: (B, H, nb, top_k, bs, D) → reshape to (B, H, nb, top_k*bs, D)
        (top_indices.unsqueeze(-1).unsqueeze(-1).expand(B, H, num_blocks, top_k, bs, D))
        k_sel = torch.gather(
            k_blocks.unsqueeze(3).expand(-1, -1, -1, top_k, -1, -1),
            2,
            top_indices.reshape(B, H, num_blocks, top_k, 1, 1).expand(
                B, H, num_blocks, top_k, bs, D
            ),
        )
        # Above gather is awkward; use index_select-free approach:
        # reshape k_blocks to (B,H,nb,bs*D) and gather along dim=2 with flattened indices
        k_flat = k_blocks.reshape(B, H, num_blocks, bs * D)
        v_blocks.reshape(B, H, num_blocks, bs * D)

        # Build absolute gather indices
        idx_flat = top_indices.reshape(B, H, num_blocks, top_k, 1).expand(-1, -1, -1, -1, bs)
        idx_flat = idx_flat * bs + torch.arange(bs, device=device).view(1, 1, 1, 1, bs)
        idx_flat = idx_flat.reshape(B, H, num_blocks, top_k * bs)

        k_sel = torch.gather(
            k_flat.unsqueeze(3).expand(-1, -1, -1, top_k, -1),
            -1,
            idx_flat.unsqueeze(-1).expand(-1, -1, -1, -1, D),
        )  # (B,H,nb,top_k,D)
        # Actually this is getting confusing. Simpler:
        # flatten indices to a per-(B,H,qblock) gather along last dim of (B,H,nb,bs*D).
        k_sel = torch.empty(B, H, num_blocks, top_k * bs, D, device=device, dtype=k_blocks.dtype)
        v_sel = torch.empty_like(k_sel)
        valid_sel = torch.empty(B, H, num_blocks, top_k * bs, device=device, dtype=torch.bool)

        for j in range(top_k):
            bj = top_indices[..., j]  # (B, H, nb)
            bj_flat_idx = bj.view(B, H, num_blocks, 1, 1).expand(-1, -1, -1, bs, D)
            k_sel[:, :, :, j * bs : (j + 1) * bs, :] = torch.gather(k_blocks, 2, bj_flat_idx)
            v_sel[:, :, :, j * bs : (j + 1) * bs, :] = torch.gather(v_blocks, 2, bj_flat_idx)
            bj_flat_idx_v = bj.view(B, H, num_blocks, 1).expand(-1, -1, -1, bs)
            valid_sel[:, :, :, j * bs : (j + 1) * bs] = torch.gather(valid_blocks, 2, bj_flat_idx_v)

        # ---- attention within selected blocks ----
        q_all = q_pad.reshape(B, H, num_blocks, bs, D)
        scale = D**-0.5
        scores = torch.matmul(q_all.float(), k_sel.float().transpose(-2, -1).contiguous()) * scale

        # FA-FIX-03: causal mask via ABSOLUTE positions (per-query position vs per-key absolute position)
        if is_causal:
            q_abs = torch.arange(num_blocks, device=device).view(num_blocks, 1) * bs + torch.arange(
                bs, device=device
            ).view(1, bs)  # (nb, bs)
            q_abs = q_abs.view(1, 1, num_blocks, bs)  # (1, 1, nb, bs)
            # Absolute key positions for each query block's selected keys
            k_abs = idx_flat  # (B, H, nb, top_k*bs) — already absolute
            causal_b = k_abs.unsqueeze(-2) <= q_abs.unsqueeze(-1)  # (B, H, nb, bs, top_k*bs)
            scores = scores.masked_fill(~causal_b, float("-inf"))

        # FA-FIX-04: padding mask on the gathered keys
        scores = scores.masked_fill(~valid_sel.unsqueeze(-2), float("-inf"))

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(v_sel.dtype)
        probs = torch.nan_to_num(probs, nan=0.0)
        probs = self.dropout(probs)

        # FA-FIX-09: head_mask applied here
        if head_mask is not None:
            hm = head_mask.to(probs.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, H, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], H, 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            probs = probs * hm

        out_blocks = torch.matmul(probs, v_sel)  # (B, H, nb, bs, D)

        # FA-FIX-08: NaN guard
        out_blocks = torch.nan_to_num(out_blocks, nan=0.0, posinf=1e4, neginf=-1e4)

        context = out_blocks.reshape(B, H, N_pad, D)[:, :, :N, :]
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        # FA-FIX-10: return block attention weights when requested
        if output_attentions:
            return context, probs
        return context, None

    # ------------------------------------------------------------------
    # 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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        device = hidden_states.device
        is_cross_attention = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross_attention else hidden_states
        kv_len_new = kv_input.shape[1]

        # ---- Q / K / V 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, kv_len_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v_new = (
            self.v_proj(kv_input)
            .view(B, kv_len_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- RoPE with past_len offset ----
        if not is_cross_attention:
            past_k_probe, _ = self._read_past_kv(past_key_value, self.layer_idx)
            past_len = 0 if past_k_probe is None else past_k_probe.shape[-2]

            if position_embeddings is not None and isinstance(position_embeddings, tuple):
                cos, sin = position_embeddings
                if past_len > 0:
                    cos = cos[..., past_len : past_len + N, :]
                    sin = sin[..., past_len : past_len + N, :]
                q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin, position_ids)
            elif self.rotary_emb is not None:
                # FA-FIX-07: use hidden_states for shape probing (was `v`)
                try:
                    cos, sin = self.rotary_emb(hidden_states, seq_len=past_len + N)
                except TypeError:
                    try:
                        cos, sin = self.rotary_emb(hidden_states, position_ids)
                    except Exception:
                        cos, sin = None, None
                if cos is not None and sin is not None:
                    if past_len > 0:
                        cos = cos[..., past_len : past_len + N, :]
                        sin = sin[..., past_len : past_len + N, :]
                    q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin, position_ids)

        # ---- KV-cache concat + update ----
        # FA-FIX-06: honor past_key_value
        past_k, past_v = self._read_past_kv(past_key_value, self.layer_idx)
        if past_k is not None and not is_cross_attention:
            k_cat = torch.cat([past_k, k_new], dim=-2)
            v_cat = torch.cat([past_v, v_new], dim=-2)
        else:
            k_cat, v_cat = k_new, v_new

        if (
            past_key_value is not None
            and hasattr(past_key_value, "update")
            and not is_cross_attention
        ):
            past_key_value.update(k_new, v_new, self.layer_idx)

        # ---- GQA ----
        if self.num_kv_heads < self.num_heads:
            k_cat = repeat_kv(k_cat, self.kv_repeat)
            v_cat = repeat_kv(v_cat, self.kv_repeat)

        N_kv = k_cat.shape[2]

        # ---- masks ----
        eff_mask = encoder_attention_mask if is_cross_attention else attention_mask
        valid_mask = self._normalize_valid_mask(eff_mask, B, N_kv, device)

        # FA-FIX-03 / FA-FIX-18: causality condition
        is_causal = self.is_decoder and not is_cross_attention and N > 1

        # ---- dispatch ----
        if N <= 128 or is_cross_attention:
            context, attn_weights = self._dense_attention(
                q,
                k_cat,
                v_cat,
                valid_mask,
                head_mask,
                is_causal=is_causal and N_kv == N,
                output_attentions=bool(output_attentions),
            )
        else:
            context, attn_weights = self._block_sparse_attention(
                q,
                k_cat,
                v_cat,
                valid_mask,
                head_mask,
                is_causal=is_causal and N_kv == N,
                output_attentions=bool(output_attentions),
            )

        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return out, (attn_weights if output_attentions else None)

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.tensor(0.0, device=self.out_proj.weight.device)

Hybrid sparse attention with block gating.

Two execution paths: * Short sequence (N <= 128) or cross-attention → dense attention * Long sequence → block-gated sparse

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.

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_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 get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.tensor(0.0, device=self.out_proj.weight.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class FactorisedLinear (in_features: int,
out_features: int,
rank: int = 64,
bias: bool = True,
energy: float = 0.95,
lora_rank: int = 16,
lora_alpha: float = 32.0,
device: torch.device | None = None,
dtype: torch.dtype | None = None)

Expand source code

class FactorisedLinear(nn.Module):
    is_custom_attention: bool = True
    r"""SVD-Factorized Linear Transformation for NOVA Architecture."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 64,
        bias: bool = True,
        energy: float = 0.95,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = min(rank, in_features, out_features)
        self.energy = energy
        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.down: nn.Linear = nn.Linear(
            in_features, self.rank, bias=False, device=device, dtype=dtype
        )
        self.up: nn.Linear = nn.Linear(
            self.rank, out_features, bias=bias, device=device, dtype=dtype
        )

        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(
                torch.randn(lora_rank, in_features, device=device, dtype=dtype) * 0.01
            )
            self.lora_B = nn.Parameter(
                torch.zeros(out_features, lora_rank, device=device, dtype=dtype)
            )
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.orthogonal_(self.down.weight)
        nn.init.orthogonal_(self.up.weight)
        if self.up.bias is not None:
            nn.init.zeros_(self.up.bias)
        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 decompose_from_weight(
        self,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
        energy: float | None = None,
    ) -> None:
        with torch.no_grad():
            W = weight.float()
            # Handle KV dimension matching for GQA
            if W.shape[0] < self.out_features and self.out_features % W.shape[0] == 0:
                rep = self.out_features // W.shape[0]
                W = W.repeat(rep, 1)
                if bias is not None:
                    bias = bias.repeat(rep)
            elif W.shape[0] > self.out_features and W.shape[0] % self.out_features == 0:
                rep = W.shape[0] // self.out_features
                W = W.view(self.out_features, rep, -1).mean(dim=1)
                if bias is not None:
                    bias = bias.view(self.out_features, rep).mean(dim=1)
            elif W.shape[0] != self.out_features or W.shape[1] != self.in_features:
                logger.warning(
                    f"FactorisedLinear shape mismatch: W={W.shape}, expected=({self.out_features}, {self.in_features}). "
                    f"Padding or truncating as fallback."
                )
                padded_W = torch.zeros(
                    self.out_features, self.in_features, device=W.device, dtype=W.dtype
                )
                padded_W[
                    : min(self.out_features, W.shape[0]),
                    : min(self.in_features, W.shape[1]),
                ] = W[
                    : min(self.out_features, W.shape[0]),
                    : min(self.in_features, W.shape[1]),
                ]
                W = padded_W

            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            target_energy = energy if energy is not None else self.energy
            total_energy = (S**2).sum()
            cum_energy = torch.cumsum(S**2, dim=0) / (total_energy + 1e-12)
            calc_rank = int((cum_energy < target_energy).sum().item()) + 1
            r = min(self.rank, calc_rank, U.size(1), Vt.size(0))
            r = max(1, r)

            sqrt_s = torch.sqrt(S[:r])
            A = (U[:, :r] * sqrt_s).to(self.up.weight.dtype).to(self.up.weight.device)
            B = (
                (sqrt_s.unsqueeze(1) * Vt[:r, :])
                .to(self.down.weight.dtype)
                .to(self.down.weight.device)
            )

            self.up.weight.data.zero_()
            self.down.weight.data.zero_()
            self.up.weight.data[:, :r] = A
            self.down.weight.data[:r, :] = B

            if bias is not None and self.up.bias is not None:
                self.up.bias.data[: min(len(self.up.bias), len(bias))].copy_(
                    bias[: min(len(self.up.bias), len(bias))].to(
                        self.up.bias.device, self.up.bias.dtype
                    )
                )

            self.freeze_base_weights()

    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_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

    @property
    def weight(self) -> torch.Tensor:
        base_w = torch.matmul(self.up.weight, self.down.weight)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            return base_w + (self.lora_B @ self.lora_A) * self.lora_scaling
        return base_w

    @property
    def bias(self) -> torch.Tensor | None:
        return self.up.bias

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.down.weight.device != x.device:
            self.to(device=x.device)
        out = self.up(self.down(x))
        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.to(dtype=x.dtype)),
                    self.lora_B.to(dtype=x.dtype),
                )
                * self.lora_scaling
            )
        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

SVD-Factorized Linear Transformation for NOVA Architecture.

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:
    base_w = torch.matmul(self.up.weight, self.down.weight)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        return base_w + (self.lora_B @ self.lora_A) * self.lora_scaling
    return base_w

Methods

def decompose_from_weight(self,
weight: torch.Tensor,
bias: torch.Tensor | None = None,
energy: float | None = None) ‑> None

Expand source code

def decompose_from_weight(
    self,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None,
    energy: float | None = None,
) -> None:
    with torch.no_grad():
        W = weight.float()
        # Handle KV dimension matching for GQA
        if W.shape[0] < self.out_features and self.out_features % W.shape[0] == 0:
            rep = self.out_features // W.shape[0]
            W = W.repeat(rep, 1)
            if bias is not None:
                bias = bias.repeat(rep)
        elif W.shape[0] > self.out_features and W.shape[0] % self.out_features == 0:
            rep = W.shape[0] // self.out_features
            W = W.view(self.out_features, rep, -1).mean(dim=1)
            if bias is not None:
                bias = bias.view(self.out_features, rep).mean(dim=1)
        elif W.shape[0] != self.out_features or W.shape[1] != self.in_features:
            logger.warning(
                f"FactorisedLinear shape mismatch: W={W.shape}, expected=({self.out_features}, {self.in_features}). "
                f"Padding or truncating as fallback."
            )
            padded_W = torch.zeros(
                self.out_features, self.in_features, device=W.device, dtype=W.dtype
            )
            padded_W[
                : min(self.out_features, W.shape[0]),
                : min(self.in_features, W.shape[1]),
            ] = W[
                : min(self.out_features, W.shape[0]),
                : min(self.in_features, W.shape[1]),
            ]
            W = padded_W

        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        target_energy = energy if energy is not None else self.energy
        total_energy = (S**2).sum()
        cum_energy = torch.cumsum(S**2, dim=0) / (total_energy + 1e-12)
        calc_rank = int((cum_energy < target_energy).sum().item()) + 1
        r = min(self.rank, calc_rank, U.size(1), Vt.size(0))
        r = max(1, r)

        sqrt_s = torch.sqrt(S[:r])
        A = (U[:, :r] * sqrt_s).to(self.up.weight.dtype).to(self.up.weight.device)
        B = (
            (sqrt_s.unsqueeze(1) * Vt[:r, :])
            .to(self.down.weight.dtype)
            .to(self.down.weight.device)
        )

        self.up.weight.data.zero_()
        self.down.weight.data.zero_()
        self.up.weight.data[:, :r] = A
        self.down.weight.data[:r, :] = B

        if bias is not None and self.up.bias is not None:
            self.up.bias.data[: min(len(self.up.bias), len(bias))].copy_(
                bias[: min(len(self.up.bias), len(bias))].to(
                    self.up.bias.device, self.up.bias.dtype
                )
            )

        self.freeze_base_weights()

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self.down.weight.device != x.device:
        self.to(device=x.device)
    out = self.up(self.down(x))
    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.to(dtype=x.dtype)),
                self.lora_B.to(dtype=x.dtype),
            )
            * self.lora_scaling
        )
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    self.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:
        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.orthogonal_(self.down.weight)
    nn.init.orthogonal_(self.up.weight)
    if self.up.bias is not None:
        nn.init.zeros_(self.up.bias)
    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 FactorisedLinear_CUR (in_features: int,
out_features: int,
cur_rank: int = 16,
bias: bool = True,
lora_rank: int = 16,
lora_alpha: float = 32.0,
freeze_base: bool = False,
device: torch.device | None = None,
dtype: torch.dtype | None = None)

Expand source code

class FactorisedLinear_CUR(nn.Module):
    is_custom_attention: bool = True
    r"""Linear projection compressed via leverage-score CUR matrix decomposition with LoRA support."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        cur_rank: int = 16,
        bias: bool = True,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        freeze_base: bool = False,
        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.cur_rank = max(1, min(cur_rank, in_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.down = nn.Linear(in_features, self.cur_rank, bias=False, **factory_kwargs)
        self.up = nn.Linear(self.cur_rank, out_features, bias=bias, **factory_kwargs)

        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.down.weight, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.up.weight, a=math.sqrt(5))
        if self.up.bias is not None:
            nn.init.zeros_(self.up.bias)

        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.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:
            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 copy_cur(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
        with torch.no_grad():
            W = weight.float()
            out_dim, in_dim = W.shape
            rank = min(self.cur_rank, out_dim, in_dim)

            try:
                U_full, S_full, Vt_full = torch.linalg.svd(W, full_matrices=False)
                col_leverage = (Vt_full**2).sum(dim=0)
                col_leverage = col_leverage / torch.clamp(col_leverage.sum(), min=1e-5)
                row_leverage = (U_full**2).sum(dim=1)
                row_leverage = row_leverage / torch.clamp(row_leverage.sum(), min=1e-5)

                _, col_idx = torch.topk(col_leverage, k=rank, largest=True)
                _, row_idx = torch.topk(row_leverage, k=rank, largest=True)

                C = W[:, col_idx]
                R = W[row_idx, :]
                W_intersection = W[row_idx][:, col_idx]
                U_int = torch.linalg.pinv(W_intersection, rcond=1e-3)
                C_Uint = C @ U_int
                rec_norm = float(torch.linalg.norm(C_Uint @ R))
                orig_norm = float(torch.linalg.norm(W))

                if rec_norm > 2.0 * orig_norm or not math.isfinite(rec_norm):
                    # Ill-conditioned CUR intersection -> fallback to stable truncated SVD
                    U_r = U_full[:, :rank]
                    S_r = S_full[:rank]
                    Vt_r = Vt_full[:rank, :]
                    sqrt_s = torch.sqrt(S_r)
                    self.down.weight.data.zero_()
                    self.down.weight.data[:rank, :in_dim].copy_(
                        (sqrt_s.unsqueeze(1) * Vt_r).to(dtype=self.down.weight.dtype)
                    )
                    self.up.weight.data.zero_()
                    self.up.weight.data[:out_dim, :rank].copy_(
                        (U_r * sqrt_s.unsqueeze(0)).to(dtype=self.up.weight.dtype)
                    )
                else:
                    self.down.weight.data.zero_()
                    self.down.weight.data[:rank, :in_dim].copy_(R.to(dtype=self.down.weight.dtype))
                    self.up.weight.data.zero_()
                    self.up.weight.data[:out_dim, :rank].copy_(
                        C_Uint.to(dtype=self.up.weight.dtype)
                    )
            except Exception:
                # SVD fallback
                nn.init.kaiming_uniform_(self.down.weight, a=math.sqrt(5))
                nn.init.kaiming_uniform_(self.up.weight, a=math.sqrt(5))

            if bias is not None and self.up.bias is not None:
                self.up.bias.data.copy_(bias.to(dtype=self.up.bias.dtype))

        if self.freeze_base:
            self.freeze_base_weights()

    @property
    def weight(self) -> torch.Tensor:
        return torch.matmul(self.up.weight, self.down.weight)

    @property
    def bias(self) -> torch.Tensor | None:
        return self.up.bias

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

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

Linear projection compressed via leverage-score CUR matrix decomposition with LoRA support.

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:
    return torch.matmul(self.up.weight, self.down.weight)

Methods

def copy_cur(self, weight: torch.Tensor, bias: torch.Tensor | None = None) ‑> None

Expand source code

def copy_cur(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
    with torch.no_grad():
        W = weight.float()
        out_dim, in_dim = W.shape
        rank = min(self.cur_rank, out_dim, in_dim)

        try:
            U_full, S_full, Vt_full = torch.linalg.svd(W, full_matrices=False)
            col_leverage = (Vt_full**2).sum(dim=0)
            col_leverage = col_leverage / torch.clamp(col_leverage.sum(), min=1e-5)
            row_leverage = (U_full**2).sum(dim=1)
            row_leverage = row_leverage / torch.clamp(row_leverage.sum(), min=1e-5)

            _, col_idx = torch.topk(col_leverage, k=rank, largest=True)
            _, row_idx = torch.topk(row_leverage, k=rank, largest=True)

            C = W[:, col_idx]
            R = W[row_idx, :]
            W_intersection = W[row_idx][:, col_idx]
            U_int = torch.linalg.pinv(W_intersection, rcond=1e-3)
            C_Uint = C @ U_int
            rec_norm = float(torch.linalg.norm(C_Uint @ R))
            orig_norm = float(torch.linalg.norm(W))

            if rec_norm > 2.0 * orig_norm or not math.isfinite(rec_norm):
                # Ill-conditioned CUR intersection -> fallback to stable truncated SVD
                U_r = U_full[:, :rank]
                S_r = S_full[:rank]
                Vt_r = Vt_full[:rank, :]
                sqrt_s = torch.sqrt(S_r)
                self.down.weight.data.zero_()
                self.down.weight.data[:rank, :in_dim].copy_(
                    (sqrt_s.unsqueeze(1) * Vt_r).to(dtype=self.down.weight.dtype)
                )
                self.up.weight.data.zero_()
                self.up.weight.data[:out_dim, :rank].copy_(
                    (U_r * sqrt_s.unsqueeze(0)).to(dtype=self.up.weight.dtype)
                )
            else:
                self.down.weight.data.zero_()
                self.down.weight.data[:rank, :in_dim].copy_(R.to(dtype=self.down.weight.dtype))
                self.up.weight.data.zero_()
                self.up.weight.data[:out_dim, :rank].copy_(
                    C_Uint.to(dtype=self.up.weight.dtype)
                )
        except Exception:
            # SVD fallback
            nn.init.kaiming_uniform_(self.down.weight, a=math.sqrt(5))
            nn.init.kaiming_uniform_(self.up.weight, a=math.sqrt(5))

        if bias is not None and self.up.bias is not None:
            self.up.bias.data.copy_(bias.to(dtype=self.up.bias.dtype))

    if self.freeze_base:
        self.freeze_base_weights()

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_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
        out = out + lora_out
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    self.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:
        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.down.weight, a=math.sqrt(5))
    nn.init.kaiming_uniform_(self.up.weight, a=math.sqrt(5))
    if self.up.bias is not None:
        nn.init.zeros_(self.up.bias)

    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 FactorisedLoRALinear (in_features: int,
out_features: int,
svd_rank: int,
lora_rank: int = 0,
lora_alpha: float = 32.0,
bias: bool = True,
dtype: torch.dtype | None = None,
device: torch.device | None = None)

Expand source code

class FactorisedLoRALinear(nn.Module):
    """
    Low-rank factorized linear with optional LoRA:

        y = (x V S U^T) + (x lora_A^T lora_B^T) * scaling + b

    Base factors (U, S, V) are frozen parameters (visible in .parameters()).
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        in_features: int,
        out_features: int,
        svd_rank: int,
        lora_rank: int = 0,
        lora_alpha: float = 32.0,
        bias: bool = True,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ):
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.svd_rank = max(1, min(svd_rank, self.in_features, self.out_features))
        self.lora_rank = max(0, int(lora_rank))
        self.lora_scaling = (lora_alpha / self.lora_rank) if self.lora_rank > 0 else 1.0

        fk = {"dtype": dtype, "device": device}
        self.U = nn.Parameter(
            torch.zeros(out_features, self.svd_rank, **fk),
            requires_grad=False,
        )
        self.S = nn.Parameter(
            torch.ones(self.svd_rank, **fk),
            requires_grad=False,
        )
        self.V = nn.Parameter(
            torch.zeros(in_features, self.svd_rank, **fk),
            requires_grad=False,
        )

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, **fk), requires_grad=False)
        else:
            self.register_parameter("bias", None)

        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty(self.lora_rank, in_features, **fk))
            self.lora_B = nn.Parameter(torch.zeros(out_features, self.lora_rank, **fk))
            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)

    def set_svd_factors(self, W: torch.Tensor, bias: torch.Tensor | None = None) -> None:
        with torch.no_grad():
            W_f = W.float()
            if W_f.shape != (self.out_features, self.in_features):
                raise ValueError(
                    f"FactorisedLoRALinear: shape mismatch "
                    f"{tuple(W_f.shape)} vs ({self.out_features}, {self.in_features})"
                )
            U, S, Vt = torch.linalg.svd(W_f, full_matrices=False)
            r = min(self.svd_rank, S.shape[0])
            self.U.data.zero_()
            self.S.data.zero_()
            self.V.data.zero_()
            self.U.data[:, :r].copy_(U[:, :r].to(device=self.U.device, dtype=self.U.dtype))
            self.S.data[:r].copy_(S[:r].to(device=self.S.device, dtype=self.S.dtype))
            self.V.data[:, :r].copy_(
                Vt[:r, :].t().contiguous().to(device=self.V.device, dtype=self.V.dtype)
            )
            if bias is not None and self.bias is not None:
                n = min(self.bias.numel(), bias.numel())
                self.bias.data[:n].copy_(
                    bias[:n].to(device=self.bias.device, dtype=self.bias.dtype)
                )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_dtype = x.dtype
        x_f = x.float()
        xV = torch.matmul(x_f, self.V.float())
        xV_S = xV * self.S.float().unsqueeze(0).unsqueeze(0)
        out = torch.matmul(xV_S, self.U.float().t().contiguous())
        if self.bias is not None:
            out = out + self.bias.float()
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            xa = torch.matmul(x_f, self.lora_A.float().t().contiguous())
            out = out + torch.matmul(xa, self.lora_B.float().t().contiguous()) * self.lora_scaling
        return out.to(orig_dtype)

    @property
    def weight(self) -> torch.Tensor:
        W = (self.U.float() * self.S.float().unsqueeze(0)) @ self.V.float().t().contiguous()
        return W

    def freeze_base_weights(self) -> None:
        self.U.requires_grad = False
        self.S.requires_grad = False
        self.V.requires_grad = False
        if self.bias is not None:
            self.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

Low-rank factorized linear with optional LoRA:

y = (x V S U^T) + (x lora_A^T lora_B^T) * scaling + b

Base factors (U, S, V) are frozen parameters (visible in .parameters()).

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

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    W = (self.U.float() * self.S.float().unsqueeze(0)) @ self.V.float().t().contiguous()
    return W

Methods

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_dtype = x.dtype
    x_f = x.float()
    xV = torch.matmul(x_f, self.V.float())
    xV_S = xV * self.S.float().unsqueeze(0).unsqueeze(0)
    out = torch.matmul(xV_S, self.U.float().t().contiguous())
    if self.bias is not None:
        out = out + self.bias.float()
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        xa = torch.matmul(x_f, self.lora_A.float().t().contiguous())
        out = out + torch.matmul(xa, self.lora_B.float().t().contiguous()) * self.lora_scaling
    return out.to(orig_dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    self.U.requires_grad = False
    self.S.requires_grad = False
    self.V.requires_grad = False
    if self.bias is not None:
        self.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 set_svd_factors(self, W: torch.Tensor, bias: torch.Tensor | None = None) ‑> None

Expand source code

def set_svd_factors(self, W: torch.Tensor, bias: torch.Tensor | None = None) -> None:
    with torch.no_grad():
        W_f = W.float()
        if W_f.shape != (self.out_features, self.in_features):
            raise ValueError(
                f"FactorisedLoRALinear: shape mismatch "
                f"{tuple(W_f.shape)} vs ({self.out_features}, {self.in_features})"
            )
        U, S, Vt = torch.linalg.svd(W_f, full_matrices=False)
        r = min(self.svd_rank, S.shape[0])
        self.U.data.zero_()
        self.S.data.zero_()
        self.V.data.zero_()
        self.U.data[:, :r].copy_(U[:, :r].to(device=self.U.device, dtype=self.U.dtype))
        self.S.data[:r].copy_(S[:r].to(device=self.S.device, dtype=self.S.dtype))
        self.V.data[:, :r].copy_(
            Vt[:r, :].t().contiguous().to(device=self.V.device, dtype=self.V.dtype)
        )
        if bias is not None and self.bias is not None:
            n = min(self.bias.numel(), bias.numel())
            self.bias.data[:n].copy_(
                bias[:n].to(device=self.bias.device, dtype=self.bias.dtype)
            )

class FactorisedMatrixState (config: Any | None = None,
rank: int | None = 32,
d_model: int | Any | None = None,
num_heads: int | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "fms", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "factorised_matrix_state",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class FactorisedMatrixState(BaseAttention):
    r"""Bidirectional & Causal Factorised Matrix State Recurrent Linear Attention."""

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 32,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        if d_model is None:
            if config is not None:
                d_model = getattr(config, "hidden_size", getattr(config, "dim", 768))
            elif source_attn is not None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                d_model = getattr(out_layer, "out_features", 768) if out_layer is not None else 768
            else:
                d_model = kwargs.get("hidden_size", 768)

        if hasattr(d_model, "hidden_size"):
            d_model = d_model.hidden_size

        if num_heads is None:
            if config is not None:
                num_heads = getattr(config, "num_attention_heads", getattr(config, "num_heads", 12))
            elif source_attn is not None:
                num_heads = getattr(
                    source_attn,
                    "n_heads",
                    getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(source_attn, "num_heads", 12),
                    ),
                )
            else:
                num_heads = kwargs.get("num_attention_heads", 12)

        super().__init__(config, rank)
        self.d_model: int = int(d_model)
        self.num_heads: int = int(num_heads)
        self.num_kv_heads: int = (
            getattr(config, "num_key_value_heads", self.num_heads)
            if config is not None
            else self.num_heads
        )
        self.head_dim: int = getattr(config, "head_dim", None) or (
            self.d_model // max(1, self.num_heads)
        )
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.proj_rank: int = rank if rank is not None else 32

        # Static layer index
        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx: int = int(resolved_layer_id)

        # LoRA adaptation configuration
        if config is not None:
            self.lora_rank: int = getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            self.lora_alpha: float = getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        elif comp_config is not None:
            self.lora_rank: int = getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16))
            self.lora_alpha: float = getattr(
                comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)
            )
        else:
            self.lora_rank: int = kwargs.get("lora_rank", 16)
            self.lora_alpha: float = kwargs.get("lora_alpha", 32.0)
        self.freeze_base: bool = self.lora_rank > 0

        # Decoder & nested encoder detection
        model_type = getattr(config, "model_type", "").lower()
        is_dec = getattr(
            config,
            "is_decoder",
            getattr(config, "is_causal", kwargs.get("is_causal", False)),
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(_mt.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)
        self.is_nested_encoder: bool = False

        # Multi-head projections equipped with LoRA adapters
        self.q_proj = LoRALinear(
            self.d_model,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.d_model,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.d_model,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Factorized low-rank projection parameters L and R
        # L is initialized with Kaiming uniform and R with zeros so that the initial low-rank perturbation
        # is exactly zero (W_res alone acts as base projection), but gradients into L are non-zero.
        self.L = nn.Parameter(torch.empty(self.all_head_size, self.proj_rank))
        self.R = nn.Parameter(torch.zeros(self.d_model, self.proj_rank))
        nn.init.kaiming_uniform_(self.L, a=math.sqrt(5))
        self.W_res = LoRALinear(
            self.all_head_size,
            self.d_model,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.beta = nn.Parameter(torch.ones(self.num_heads, 1, 1) * 2.0)
        self.rotary_emb = None
        self.comp_config: Any | None = comp_config

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    def copy_weights(self, old_attn: Any) -> None:
        """Transfer pretrained projection weights with strict Frobenius norm verification."""
        if old_attn is None:
            return

        cls_name = old_attn.__class__.__name__
        has_out_lin = (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        )
        is_nested = False
        if (
            (
                hasattr(old_attn, "self")
                and hasattr(old_attn.self, "query")
                and hasattr(old_attn, "output")
                and hasattr(old_attn.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)

        if hasattr(old_attn, "rotary_emb"):
            self.rotary_emb = old_attn.rotary_emb

        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 + c_proj
        if hasattr(old_attn, "c_attn") and old_attn.c_attn is not None:
            c_attn_mod = old_attn.c_attn
            if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    print(f"copy_weights failed: {e}")
                    logger.warning(f"copy_weights failed: {e}")
            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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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(old_attn, "c_proj") and old_attn.c_proj is not None:
                c_proj_mod = old_attn.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()
                    if (
                        (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])
                    ):
                        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. RoBERTa / BERT nested self-attention
        elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
            s = old_attn.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(old_attn, "output") and hasattr(old_attn.output, "dense"):
                out_mod = old_attn.output.dense
                extracted["o"] = (
                    out_mod.weight.data.clone(),
                    out_mod.bias.data.clone() if out_mod.bias is not None else None,
                )

        # 3. Standard linear projections (DistilBERT, Qwen, LLaMA)
        else:
            proj_candidates = {
                "q": ["q_lin", "q_proj", "query", "W_q"],
                "k": ["k_lin", "k_proj", "key", "W_k"],
                "v": ["v_lin", "v_proj", "value", "W_v"],
                "o": ["out_lin", "out_proj", "dense", "o_proj"],
            }
            for proj_key, names in proj_candidates.items():
                for name in names:
                    if hasattr(old_attn, name):
                        mod = getattr(old_attn, 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

        # Fallback via UniversalAdapter
        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(old_attn, 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"FactorisedMatrixState copy_weights: missing mandatory projections {missing_after} "
                f"from module {old_attn.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.W_res,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

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

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    def _matrix_scan_linear(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        decay: torch.Tensor,
        mask: torch.Tensor | None = None,
        reverse: bool = False,
        chunk_size: int = 64,
    ) -> torch.Tensor:
        r"""
        Chunked associative matrix scan with non-negative ELU+1 feature maps.
        Optimized to prevent VRAM explosions during outer-product state accumulation.
        """
        B, H, L, d = q.shape

        if HAS_TRITON_FMS and q.is_cuda and d <= MAX_REGISTER_DIM and mask is None and not reverse:
            try:
                decay_scalar = decay.float().clamp(min=1e-4, max=1.0 - 1e-4).squeeze()
                if decay_scalar.dim() == 0:
                    decay_scalar = decay_scalar.repeat(H)
                return fms_scan_triton(q, k, v, decay_scalar).to(q.dtype)
            except Exception as e:
                logger.debug("FMS Triton failed, falling back: %s", e)

        q_f = F.elu(q.float()) + 1.0
        k_f = (F.elu(k.float()) + 1.0) / math.sqrt(d)
        v_f = v.float()
        decay_f = decay.float().clamp(min=1e-4, max=1.0 - 1e-4).view(1, H, 1, 1)

        if mask is not None:
            m_f = mask.float().unsqueeze(1) if mask.dim() == 2 else mask.float()
            if m_f.dim() == 4:
                m_f = (mask > -1000.0).any(dim=1).any(dim=-2).float().unsqueeze(1)
            k_f = k_f * m_f.unsqueeze(-1)
            v_f = v_f * m_f.unsqueeze(-1)

        if reverse:
            q_f = q_f.flip(dims=[2])
            k_f = k_f.flip(dims=[2])
            v_f = v_f.flip(dims=[2])

        if chunk_size >= L:
            log_decay = torch.log(decay_f)
            idx = torch.arange(L, device=q.device, dtype=torch.float32)
            dist = idx.unsqueeze(1) - idx.unsqueeze(0)
            causal_mask = dist >= 0
            # Memory optimization: compute D without extra unsqueeze where possible
            D = torch.exp(dist.unsqueeze(0).unsqueeze(0) * log_decay).masked_fill(
                ~causal_mask.unsqueeze(0).unsqueeze(0), 0.0
            )

            # Avoid .contiguous() after transpose if not absolutely required
            scores = torch.matmul(q_f, k_f.transpose(-1, -2)) * D
            num = torch.matmul(scores, v_f)
            denom = scores.sum(dim=-1, keepdim=True).clamp_min(1e-5)
            y = num / denom
        else:
            num_chunks = math.ceil(L / chunk_size)
            y_chunks = []

            S = torch.zeros(B, H, d, d, device=q.device, dtype=torch.float32)
            z = torch.zeros(B, H, d, 1, device=q.device, dtype=torch.float32)

            log_decay = torch.log(decay_f)

            for i in range(num_chunks):
                start = i * chunk_size
                end = min(start + chunk_size, L)
                curr_c = end - start

                qc = q_f[:, :, start:end, :]
                kc = k_f[:, :, start:end, :]
                vc = v_f[:, :, start:end, :]

                idx_c = torch.arange(curr_c, device=q.device, dtype=torch.float32)
                dist_c = idx_c.unsqueeze(1) - idx_c.unsqueeze(0)
                causal_mask_c = dist_c >= 0
                Dc = torch.exp(dist_c.unsqueeze(0).unsqueeze(0) * log_decay).masked_fill(
                    ~causal_mask_c.unsqueeze(0).unsqueeze(0), 0.0
                )

                # Intra-chunk
                scores_c = torch.matmul(qc, kc.transpose(-1, -2)) * Dc
                num_intra = torch.matmul(scores_c, vc)
                denom_intra = scores_c.sum(dim=-1, keepdim=True)

                # Inter-chunk
                decay_from_prev = torch.exp(
                    (idx_c + 1).view(1, 1, curr_c, 1) * log_decay
                )  # (1, H, curr_c, 1)

                # MEMORY FIX: Use einsum instead of expanding S to (B, H, curr_c, d, d)
                # num_inter = (qc @ S) * decay_from_prev
                # qc: (B, H, C, d), S: (B, H, d, d) -> output: (B, H, C, d)
                num_inter = torch.einsum("bhcd,bhde->bhce", qc, S)
                num_inter = num_inter * decay_from_prev

                # denom_inter = (qc @ z) * decay_from_prev
                # qc: (B, H, C, d), z: (B, H, d, 1) -> output: (B, H, C, 1)
                denom_inter = torch.einsum("bhcd,bhde->bhce", qc, z)
                denom_inter = denom_inter * decay_from_prev

                num_c = num_intra + num_inter
                denom_c = (denom_intra + denom_inter).clamp_min(1e-5)
                y_chunks.append(num_c / denom_c)

                # Update state (S, z) for next chunk
                decay_chunk = torch.exp(float(curr_c) * log_decay)  # (1, H, 1, 1)

                dist_to_end = (curr_c - 1 - idx_c).view(1, 1, curr_c, 1)
                decay_to_end = torch.exp(dist_to_end * log_decay)  # (1, H, curr_c, 1)

                # MEMORY FIX: Apply decay to keys BEFORE outer product!
                # Instead of computing a huge (B,H,C,d,d) tensor, we scale kc first.
                kc_decayed = kc * decay_to_end  # (B, H, C, d)

                # Now compute the accumulated state over the chunk directly into (B,H,d,d) using einsum
                # chunk_S = sum_{c} (kc_decayed_c \otimes vc_c)
                chunk_S = torch.einsum("bhcd,bhce->bhde", kc_decayed, vc)

                # chunk_z = sum_{c} kc_decayed_c
                chunk_z = kc_decayed.sum(dim=2).unsqueeze(-1)  # (B, H, d, 1)

                S = S * decay_chunk + chunk_S
                z = z * decay_chunk + chunk_z

            y = torch.cat(y_chunks, dim=2)

        if reverse:
            y = y.flip(dims=[2])

        return y.to(q.dtype)

    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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim

        is_cross_attention = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross_attention else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Projections
        q = self.q_proj(hidden_states).view(B, L, H, d).transpose(1, 2).contiguous()
        k = self.k_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()
        v = self.v_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # 2. Rotary Position Embeddings (RoPE)
        if position_embeddings is not None and not is_cross_attention:
            if isinstance(position_embeddings, tuple):
                cos, sin = position_embeddings
            else:
                cos, sin = position_embeddings, None
            if sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        elif self.rotary_emb is not None and not is_cross_attention:
            if position_ids is not None:
                cos, sin = self.rotary_emb(v, seq_len=kv_len)
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)

        # 3. Grouped Query Attention (GQA) head expansion
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        eff_mask = encoder_attention_mask if is_cross_attention else attention_mask
        mask_bool = None
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                mask_bool = eff_mask.bool()
            elif eff_mask.dim() == 4:
                m_sub = eff_mask > -1000.0 if eff_mask.dtype != torch.bool else eff_mask
                mask_bool = m_sub.any(dim=1).any(dim=-2).bool()

        decay = torch.sigmoid(self.beta)  # [H, 1, 1]

        # 5. Associative Matrix Scan (Forward + Anticausal Backward for encoders)
        y_fwd = self._matrix_scan_linear(q, k, v, decay, mask=mask_bool, reverse=False)

        if self.is_decoder or is_cross_attention:
            y = y_fwd
        else:
            y_bwd = self._matrix_scan_linear(q, k, v, decay, mask=mask_bool, reverse=True)
            y = 0.5 * (y_fwd + y_bwd)

        # 6. Head mask handling
        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.shape[0], head_mask.shape[1], 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            y = y * head_mask

        y = torch.nan_to_num(y, nan=0.0)
        y_flat = y.transpose(1, 2).contiguous().view(B, L, self.all_head_size)

        # 7. Output projection with nested encoder bypass
        if self.is_nested_encoder:
            out = y_flat
        else:
            # Low-rank factorized bilinear output projection + residual output projection
            low_rank = torch.matmul(torch.matmul(y_flat, self.L), self.R.t().contiguous())
            out = self.W_res(y_flat) + low_rank

        return out, None

Bidirectional & Causal Factorised Matrix State Recurrent Linear Attention.

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.

Instance variables

prop W_k : LoRALinear

Expand source code

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

prop W_q : LoRALinear

Expand source code

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

prop W_v : LoRALinear

Expand source code

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

prop c_proj : LoRALinear

Expand source code

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

prop dense : LoRALinear

Expand source code

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

prop k_lin : LoRALinear

Expand source code

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

prop key : LoRALinear

Expand source code

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

prop o_proj : LoRALinear

Expand source code

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

prop out_lin : LoRALinear

Expand source code

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

prop out_proj : LoRALinear

Expand source code

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

prop q_lin : LoRALinear

Expand source code

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

prop query : LoRALinear

Expand source code

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

prop v_lin : LoRALinear

Expand source code

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

prop value : LoRALinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Transfer pretrained projection weights with strict Frobenius norm verification."""
    if old_attn is None:
        return

    cls_name = old_attn.__class__.__name__
    has_out_lin = (
        hasattr(old_attn, "out_lin")
        or hasattr(old_attn, "c_proj")
        or hasattr(old_attn, "o_proj")
    )
    is_nested = False
    if (
        (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.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)

    if hasattr(old_attn, "rotary_emb"):
        self.rotary_emb = old_attn.rotary_emb

    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 + c_proj
    if hasattr(old_attn, "c_attn") and old_attn.c_attn is not None:
        c_attn_mod = old_attn.c_attn
        if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
            try:
                c_attn_mod.merge()
            except Exception as e:
                print(f"copy_weights failed: {e}")
                logger.warning(f"copy_weights failed: {e}")
        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()
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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(old_attn, "c_proj") and old_attn.c_proj is not None:
            c_proj_mod = old_attn.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()
                if (
                    (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])
                ):
                    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. RoBERTa / BERT nested self-attention
    elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
        s = old_attn.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(old_attn, "output") and hasattr(old_attn.output, "dense"):
            out_mod = old_attn.output.dense
            extracted["o"] = (
                out_mod.weight.data.clone(),
                out_mod.bias.data.clone() if out_mod.bias is not None else None,
            )

    # 3. Standard linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q": ["q_lin", "q_proj", "query", "W_q"],
            "k": ["k_lin", "k_proj", "key", "W_k"],
            "v": ["v_lin", "v_proj", "value", "W_v"],
            "o": ["out_lin", "out_proj", "dense", "o_proj"],
        }
        for proj_key, names in proj_candidates.items():
            for name in names:
                if hasattr(old_attn, name):
                    mod = getattr(old_attn, 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

    # Fallback via UniversalAdapter
    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(old_attn, 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"FactorisedMatrixState copy_weights: missing mandatory projections {missing_after} "
            f"from module {old_attn.__class__.__name__}."
        )

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.W_res,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device, dtype=target_mod.weight.dtype
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

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

            if self.freeze_base:
                target_mod.freeze_base_weights()

Transfer pretrained projection weights with strict Frobenius norm verification.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class FactorizedLinear (in_features: int,
out_features: int,
rank: int = 128,
lora_rank: int = 32,
bias: bool = True,
device: torch.device | None = None,
dtype: torch.dtype | None = None)

Expand source code

class FactorizedLinear(nn.Module):
    is_custom_attention = True

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 128,
        lora_rank: int = 32,
        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 = (
            min(int(lora_rank), min(self.in_features, self.out_features)) if lora_rank > 0 else 0
        )

        factory_kwargs = {"device": device, "dtype": dtype}
        self.lin1 = nn.Linear(self.in_features, self.rank, bias=False, **factory_kwargs)
        self.lin2 = nn.Linear(self.rank, self.out_features, bias=bias, **factory_kwargs)

        self.lin1.weight.requires_grad = False
        self.lin2.weight.requires_grad = False
        if self.lin2.bias is not None:
            self.lin2.bias.requires_grad = False

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

        self.reset_parameters()

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

    @property
    def weight(self) -> torch.Tensor:
        w = torch.matmul(self.lin2.weight, self.lin1.weight)
        if self.lora_A is not None:
            w = w + torch.matmul(self.lora_B, self.lora_A)
        return w

    @property
    def bias(self) -> torch.Tensor | None:
        return self.lin2.bias

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_dtype = x.dtype
        x_down = x.to(self.lin1.weight.dtype)
        latent = self.lin1(x_down)
        latent_up = latent.to(self.lin2.weight.dtype)
        out = self.lin2(latent_up)
        if self.lora_A is not None:
            a = self.lora_A.to(orig_dtype) if self.lora_A.dtype != orig_dtype else self.lora_A
            b = self.lora_B.to(orig_dtype) if self.lora_B.dtype != orig_dtype else self.lora_B
            out = out.to(orig_dtype) + F.linear(F.linear(x, a), b)
        return out.to(orig_dtype)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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.lin2.bias

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    w = torch.matmul(self.lin2.weight, self.lin1.weight)
    if self.lora_A is not None:
        w = w + torch.matmul(self.lora_B, self.lora_A)
    return w

Methods

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    orig_dtype = x.dtype
    x_down = x.to(self.lin1.weight.dtype)
    latent = self.lin1(x_down)
    latent_up = latent.to(self.lin2.weight.dtype)
    out = self.lin2(latent_up)
    if self.lora_A is not None:
        a = self.lora_A.to(orig_dtype) if self.lora_A.dtype != orig_dtype else self.lora_A
        b = self.lora_B.to(orig_dtype) if self.lora_B.dtype != orig_dtype else self.lora_B
        out = out.to(orig_dtype) + F.linear(F.linear(x, a), b)
    return out.to(orig_dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def reset_parameters(self) ‑> None

Expand source code

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

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

        # RPC-16: cache
        self._cached_weight: torch.Tensor | None = None
        self._cached_versions: tuple[int, ...] = ()

    def _current_versions(self) -> tuple[int, ...]:
        v = [
            getattr(self.down.weight, "_version", 0),
            getattr(self.up.weight, "_version", 0),
        ]
        if self.lora_A is not None:
            v.append(getattr(self.lora_A, "_version", 0))
            v.append(getattr(self.lora_B, "_version", 0))
        return tuple(v)

    @property
    def weight(self) -> torch.Tensor:
        if (
            not self.training
            and self._cached_weight is not None
            and self._cached_versions == self._current_versions()
        ):
            return self._cached_weight
        w = torch.matmul(self.up.weight.float(), self.down.weight.float())
        if self.lora_A is not None and self.lora_B is not None:
            w = w + (self.lora_B.float() @ self.lora_A.float()) * self.lora_scaling
        w = w.to(self.down.weight.dtype)
        if not self.training:
            self._cached_weight = w
            self._cached_versions = self._current_versions()
        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:
        d_w = (
            self.down.weight.to(x.dtype) if self.down.weight.dtype != x.dtype else self.down.weight
        )
        u_w = self.up.weight.to(x.dtype) if self.up.weight.dtype != x.dtype else self.up.weight
        u_b = (
            self.up.bias.to(x.dtype)
            if (self.up.bias is not None and self.up.bias.dtype != x.dtype)
            else self.up.bias
        )
        out = F.linear(F.linear(x, d_w), u_w, u_b)
        if self.lora_A is not None and self.lora_B is not None:
            a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
            b = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
            out = out + F.linear(F.linear(x, a), b) * self.lora_scaling
        return out

    def storage_bytes(self) -> int:
        b = 0
        for p in self.parameters():
            b += p.numel() * p.element_size()
        return b

SVD-factorized linear + LoRA: 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:
    if (
        not self.training
        and self._cached_weight is not None
        and self._cached_versions == self._current_versions()
    ):
        return self._cached_weight
    w = torch.matmul(self.up.weight.float(), self.down.weight.float())
    if self.lora_A is not None and self.lora_B is not None:
        w = w + (self.lora_B.float() @ self.lora_A.float()) * self.lora_scaling
    w = w.to(self.down.weight.dtype)
    if not self.training:
        self._cached_weight = w
        self._cached_versions = self._current_versions()
    return w

Methods

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    d_w = (
        self.down.weight.to(x.dtype) if self.down.weight.dtype != x.dtype else self.down.weight
    )
    u_w = self.up.weight.to(x.dtype) if self.up.weight.dtype != x.dtype else self.up.weight
    u_b = (
        self.up.bias.to(x.dtype)
        if (self.up.bias is not None and self.up.bias.dtype != x.dtype)
        else self.up.bias
    )
    out = F.linear(F.linear(x, d_w), u_w, u_b)
    if self.lora_A is not None and self.lora_B is not None:
        a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
        b = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
        out = out + F.linear(F.linear(x, a), b) * self.lora_scaling
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    self.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 storage_bytes(self) ‑> int

Expand source code

def storage_bytes(self) -> int:
    b = 0
    for p in self.parameters():
        b += p.numel() * p.element_size()
    return b

class FactorizedLoRAUp (in_features: int,
out_features: int,
bias: bool = True,
lora_rank: int = 16,
lora_alpha: float = 32.0)

Expand source code

class FactorizedLoRAUp(nn.Module):
    """Up-projection with frozen base weight and trainable LoRA adapter."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.lora_rank = lora_rank
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter("bias", 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 lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty(lora_rank, in_features))
            self.lora_B = nn.Parameter(torch.empty(out_features, lora_rank))
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)
            self.scaling = lora_alpha / lora_rank
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)
            self.scaling = 1.0
        self.weight.requires_grad = False
        if self.bias is not None:
            self.bias.requires_grad = False

    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_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.scaling
        return out

Up-projection with frozen base weight and trainable LoRA adapter.

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.scaling
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def freeze_base_weights(self) ‑> None

Expand source code

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

class FlashAttention3 (config: Any | None = None,
rank: int | None = None,
use_fp8: bool = False,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "flash3", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "flash_attention3",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class FlashAttention3(BaseAttention):
    r"""FlashAttention-3 Multi-Head Exact Attention."""

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        use_fp8: bool = False,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        if config is None:

            class DummyConfig:
                pass

            config = DummyConfig()
            if source_attn is not None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                q_layer = getattr(
                    source_attn,
                    "q_lin",
                    getattr(
                        source_attn,
                        "query",
                        getattr(source_attn, "c_attn", getattr(source_attn, "q_proj", None)),
                    ),
                )
                if out_layer is not None and hasattr(out_layer, "out_features"):
                    config.hidden_size = out_layer.out_features
                elif q_layer is not None and hasattr(q_layer, "in_features"):
                    config.hidden_size = q_layer.in_features
                else:
                    config.hidden_size = kwargs.get("hidden_size", 768)
                config.num_attention_heads = getattr(
                    source_attn,
                    "n_heads",
                    getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(
                            source_attn,
                            "num_heads",
                            kwargs.get("num_heads", kwargs.get("num_attention_heads", 12)),
                        ),
                    ),
                )
            else:
                config.hidden_size = kwargs.get("hidden_size", 768)
                config.num_attention_heads = kwargs.get(
                    "num_heads", kwargs.get("num_attention_heads", 12)
                )

        super().__init__(config, rank)
        self.num_heads: int = (
            getattr(
                config,
                "num_attention_heads",
                getattr(
                    config,
                    "num_heads",
                    kwargs.get("num_heads", kwargs.get("num_attention_heads", 12)),
                ),
            )
            if config is not None
            else kwargs.get("num_heads", kwargs.get("num_attention_heads", 12))
        )
        self.num_kv_heads: int = (
            getattr(config, "num_key_value_heads", self.num_heads)
            if config is not None
            else self.num_heads
        )
        self.hidden_size: int = (
            getattr(
                config,
                "hidden_size",
                getattr(config, "d_model", kwargs.get("hidden_size", 768)),
            )
            if config is not None
            else kwargs.get("hidden_size", 768)
        )

        cfg_head_dim = getattr(config, "head_dim", None) if config is not None else None
        if cfg_head_dim is None and source_attn is not None:
            for attr in ("q_proj", "query", "q_lin", "W_q"):
                mod = getattr(source_attn, attr, None)
                if mod is not None and hasattr(mod, "weight"):
                    cfg_head_dim = mod.weight.shape[0] // max(1, self.num_heads)
                    break
        self.head_dim: int = (
            int(cfg_head_dim)
            if cfg_head_dim is not None
            else getattr(config, "d_kv", self.hidden_size // max(1, self.num_heads))
        )
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.use_fp8: bool = bool(use_fp8)

        # Static layer index
        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx: int = int(resolved_layer_id)

        # LoRA adaptation configuration
        self.lora_rank: int = getattr(
            config,
            "lora_rank",
            getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha: float = getattr(
            config,
            "lora_alpha",
            getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base: bool = self.lora_rank > 0

        # Decoder & nested encoder detection
        _mt = str(getattr(config, "model_type", "") or "").lower() if config is not None else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_dec = (
            (
                getattr(config, "is_decoder", False)
                or getattr(config, "is_causal", False)
                or any(_mt.startswith(p) for p in _decoder_prefixes)
            )
            if config is not None
            else False
        )
        if hasattr(config, "architectures") and any(
            "GPT" in str(a) or "CausalLM" in str(a)
            for a in (getattr(config, "architectures", []) or [])
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)

        is_nested = kwargs.get("is_nested_encoder")
        if is_nested is None and source_attn is not None:
            src_cls = source_attn.__class__.__name__
            if (
                src_cls
                in [
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                    "ElectraSelfAttention",
                    "DebertaSelfAttention",
                ]
                or (src_cls.endswith("SelfAttention") and "distil" not in src_cls.lower())
                or (hasattr(source_attn, "self") and hasattr(source_attn, "output"))
            ):
                is_nested = True
            elif (
                hasattr(source_attn, "out_lin")
                or hasattr(source_attn, "c_proj")
                or hasattr(source_attn, "o_proj")
            ):
                is_nested = False
        if is_nested is None:
            is_nested = False
        self.is_nested_encoder: bool = bool(is_nested)
        self.gradient_checkpointing: bool = (
            getattr(config, "gradient_checkpointing", False) if config is not None else False
        )

        # Multi-head projections equipped with LoRA adapters
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )
        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=self.freeze_base,
        )

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        self.rotary_emb = None
        self.comp_config: Any | None = comp_config

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    def copy_weights(self, src: Any) -> None:
        r"""Copy projection weights from pretrained attention module with Frobenius norm verification."""
        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 has_out_lin:
            is_nested = False
        elif (
            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())
        ):
            is_nested = True
        else:
            is_nested = False
        self.is_nested_encoder = bool(is_nested)

        if hasattr(src, "rotary_emb"):
            self.rotary_emb = src.rotary_emb

        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 + c_proj
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c_attn_mod = src.c_attn
            if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    logger.warning(f"copy_weights c_attn merge warning: {e}")
            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()
                cls_c = getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__
                base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    (cls_c == "Conv1D")
                    or (base_cls_c == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                    or (
                        c_w.ndim == 2
                        and c_w.shape[0] < c_w.shape[1]
                        and "Linear" not in cls_c
                        and "Linear" not in base_cls_c
                    )
                ):
                    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()
                    cls_o = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__
                    base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__
                    if (
                        (cls_o == "Conv1D")
                        or (base_cls_o == "Conv1D")
                        or hasattr(c_proj_mod, "nf")
                        or hasattr(base_o, "nf")
                        or (
                            o_w.ndim == 2
                            and o_w.shape[0] < o_w.shape[1]
                            and "Linear" not in cls_o
                            and "Linear" not in base_cls_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. RoBERTa / BERT nested self-attention
        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. Standard linear projections (DistilBERT, Qwen, LLaMA)
        else:
            proj_candidates = {
                "q": ["q_lin", "q_proj", "query", "W_q"],
                "k": ["k_lin", "k_proj", "key", "W_k"],
                "v": ["v_lin", "v_proj", "value", "W_v"],
                "o": ["out_lin", "out_proj", "dense", "o_proj", "proj", "output"],
            }
            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

        # Fallback via UniversalAdapter
        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(f"UniversalAdapter fallback warning: {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:
            logger.debug(
                f"FlashAttention3 copy_weights: missing mandatory projections {missing_after} "
                f"from module {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

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

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        H, Dh = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        is_cross_attention = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross_attention else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Projections
        q = self.q_proj(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
        k = (
            self.k_proj(kv_input)
            .view(B, kv_len, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(kv_input)
            .view(B, kv_len, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. Rotary Position Embeddings (RoPE)
        if position_embeddings is not None and not is_cross_attention:
            if isinstance(position_embeddings, tuple):
                cos, sin = position_embeddings
            else:
                cos, sin = position_embeddings, None
            if sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        elif self.rotary_emb is not None and not is_cross_attention:
            if position_ids is not None:
                cos, sin = self.rotary_emb(v, seq_len=kv_len)
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)

        # 3. Grouped Query Attention (GQA) head expansion
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        causal_flag = bool(self.is_decoder and not is_cross_attention and kv_len == L)

        # 4. FlashAttention-3 execution path
        if FLASH_ATTN_AVAILABLE and hidden_states.is_cuda:
            try:
                q_fa = q.transpose(1, 2).contiguous()
                k_fa = k.transpose(1, 2).contiguous()
                v_fa = v.transpose(1, 2).contiguous()

                if self.use_fp8 and hasattr(torch, "float8_e4m3fn"):
                    q_fa = q_fa.to(torch.float8_e4m3fn)
                    k_fa = k_fa.to(torch.float8_e4m3fn)
                    v_fa = v_fa.to(torch.float8_e4m3fn)

                out = flash_attn_func(q_fa, k_fa, v_fa, causal=causal_flag).to(dtype)
                out = out.contiguous().view(B, L, self.all_head_size)
            except Exception:
                out = None
        else:
            out = None

        # 5. SDPA fallback path
        if out is None:
            eff_mask = encoder_attention_mask if is_cross_attention else attention_mask
            attn_mask = None
            if eff_mask is not None:
                if eff_mask.dim() == 2:
                    attn_mask = (1.0 - eff_mask.unsqueeze(1).unsqueeze(2).float()) * -1e4
                elif eff_mask.dim() == 3:
                    attn_mask = eff_mask.unsqueeze(1).float()
                elif eff_mask.dim() == 4:
                    attn_mask = eff_mask.float()

            if causal_flag and attn_mask is None:
                out_sdpa = F.scaled_dot_product_attention(
                    q,
                    k,
                    v,
                    attn_mask=None,
                    is_causal=True,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
            else:
                if causal_flag:
                    causal_m = torch.triu(
                        torch.full(
                            (L, kv_len),
                            float("-inf"),
                            device=device,
                            dtype=torch.float32,
                        ),
                        diagonal=1,
                    )
                    attn_mask = (
                        (attn_mask + causal_m.unsqueeze(0).unsqueeze(0))
                        if attn_mask is not None
                        else causal_m.unsqueeze(0).unsqueeze(0)
                    )

                out_sdpa = F.scaled_dot_product_attention(
                    q,
                    k,
                    v,
                    attn_mask=attn_mask,
                    is_causal=False,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )

            # Head mask handling
            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.shape[0], head_mask.shape[1], 1, 1)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(-1)
                out_sdpa = out_sdpa * head_mask

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

        # 6. Output projection with nested encoder bypass
        if self.is_nested_encoder:
            final_out = out
        else:
            final_out = self.out_proj(out)

        return final_out, None

FlashAttention-3 Multi-Head Exact Attention.

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.

Instance variables

prop W_k : LoRALinear

Expand source code

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

prop W_q : LoRALinear

Expand source code

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

prop W_v : LoRALinear

Expand source code

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

prop c_proj : LoRALinear

Expand source code

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

prop dense : LoRALinear

Expand source code

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

prop k_lin : LoRALinear

Expand source code

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

prop key : LoRALinear

Expand source code

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

prop o_proj : LoRALinear

Expand source code

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

prop out_lin : LoRALinear

Expand source code

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

prop q_lin : LoRALinear

Expand source code

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

prop query : LoRALinear

Expand source code

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

prop v_lin : LoRALinear

Expand source code

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

prop value : LoRALinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    r"""Copy projection weights from pretrained attention module with Frobenius norm verification."""
    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 has_out_lin:
        is_nested = False
    elif (
        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())
    ):
        is_nested = True
    else:
        is_nested = False
    self.is_nested_encoder = bool(is_nested)

    if hasattr(src, "rotary_emb"):
        self.rotary_emb = src.rotary_emb

    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 + c_proj
    if hasattr(src, "c_attn") and src.c_attn is not None:
        c_attn_mod = src.c_attn
        if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
            try:
                c_attn_mod.merge()
            except Exception as e:
                logger.warning(f"copy_weights c_attn merge warning: {e}")
        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()
            cls_c = getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__
            base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__
            if (
                (cls_c == "Conv1D")
                or (base_cls_c == "Conv1D")
                or hasattr(c_attn_mod, "nf")
                or hasattr(base_c, "nf")
                or (
                    c_w.ndim == 2
                    and c_w.shape[0] < c_w.shape[1]
                    and "Linear" not in cls_c
                    and "Linear" not in base_cls_c
                )
            ):
                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()
                cls_o = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__
                base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    (cls_o == "Conv1D")
                    or (base_cls_o == "Conv1D")
                    or hasattr(c_proj_mod, "nf")
                    or hasattr(base_o, "nf")
                    or (
                        o_w.ndim == 2
                        and o_w.shape[0] < o_w.shape[1]
                        and "Linear" not in cls_o
                        and "Linear" not in base_cls_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. RoBERTa / BERT nested self-attention
    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. Standard linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q": ["q_lin", "q_proj", "query", "W_q"],
            "k": ["k_lin", "k_proj", "key", "W_k"],
            "v": ["v_lin", "v_proj", "value", "W_v"],
            "o": ["out_lin", "out_proj", "dense", "o_proj", "proj", "output"],
        }
        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

    # Fallback via UniversalAdapter
    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(f"UniversalAdapter fallback warning: {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:
        logger.debug(
            f"FlashAttention3 copy_weights: missing mandatory projections {missing_after} "
            f"from module {src.__class__.__name__}."
        )

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device, dtype=target_mod.weight.dtype
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

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

            if self.freeze_base:
                target_mod.freeze_base_weights()

Copy projection weights from pretrained attention module with Frobenius norm verification.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class FlashAttention3FP8 (config: Any | None = None,
rank: int | None = None,
d_model: int | Any = 768,
n_heads: int = 12,
use_fp8: bool = False,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "flash3_fp8", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "flash_attention3_fp8",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class FlashAttention3FP8(BaseAttention):
    r"""FlashAttention-3 FP8 Hardware-Accelerated Multi-Head Attention.

    Accelerates attention using FP8 (e4m3fn) precision with FlashAttention-3 kernels on
    Hopper (SM 9.0+) GPUs. Automatically falls back to native PyTorch SDPA on pre-Hopper
    hardware or CPU.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any = 768,
        n_heads: int = 12,
        use_fp8: bool = False,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        if config is None:

            class DummyConfig:
                pass

            config = DummyConfig()
            if source_attn is not None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                q_layer = getattr(
                    source_attn,
                    "q_lin",
                    getattr(
                        source_attn,
                        "query",
                        getattr(source_attn, "c_attn", getattr(source_attn, "q_proj", None)),
                    ),
                )
                if out_layer is not None and hasattr(out_layer, "out_features"):
                    config.hidden_size = out_layer.out_features
                elif q_layer is not None and hasattr(q_layer, "in_features"):
                    config.hidden_size = q_layer.in_features
                else:
                    config.hidden_size = kwargs.get("hidden_size", kwargs.get("d_model", d_model))
                config.num_attention_heads = getattr(
                    source_attn,
                    "n_heads",
                    getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(
                            source_attn,
                            "num_heads",
                            kwargs.get("num_heads", kwargs.get("num_attention_heads", n_heads)),
                        ),
                    ),
                )
            else:
                config.hidden_size = kwargs.get("hidden_size", kwargs.get("d_model", d_model))
                config.num_attention_heads = kwargs.get(
                    "num_heads", kwargs.get("num_attention_heads", n_heads)
                )

        super().__init__(config, rank)
        self.num_heads: int = (
            getattr(
                config,
                "num_attention_heads",
                getattr(
                    config,
                    "num_heads",
                    kwargs.get("num_heads", kwargs.get("num_attention_heads", n_heads)),
                ),
            )
            if config is not None
            else kwargs.get("num_heads", kwargs.get("num_attention_heads", n_heads))
        )
        self.num_kv_heads: int = kwargs.get(
            "num_kv_heads",
            kwargs.get(
                "num_key_value_heads",
                (
                    getattr(
                        config,
                        "num_key_value_heads",
                        getattr(config, "num_kv_heads", self.num_heads),
                    )
                    if config is not None
                    else self.num_heads
                ),
            ),
        )
        hidden_size: int = (
            getattr(
                config,
                "hidden_size",
                getattr(
                    config,
                    "d_model",
                    kwargs.get("hidden_size", kwargs.get("d_model", d_model)),
                ),
            )
            if config is not None
            else kwargs.get("hidden_size", kwargs.get("d_model", d_model))
        )

        cfg_head_dim = getattr(config, "head_dim", None) if config is not None else None
        if cfg_head_dim is None and source_attn is not None:
            for attr in ("q_proj", "query", "q_lin", "W_q"):
                mod = getattr(source_attn, attr, None)
                if mod is not None and hasattr(mod, "weight"):
                    cfg_head_dim = mod.weight.shape[0] // max(1, self.num_heads)
                    break
        self.head_dim: int = (
            int(cfg_head_dim)
            if cfg_head_dim is not None
            else kwargs.get("head_dim", hidden_size // max(1, self.num_heads))
        )
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))

        # Check hardware capability for FP8 (requires Hopper SM 9.0+)
        self.use_fp8: bool = bool(use_fp8)
        if self.use_fp8:
            hopper_capable = False
            if torch.cuda.is_available() and hasattr(torch.cuda, "get_device_capability"):
                try:
                    cc = torch.cuda.get_device_capability()
                    hopper_capable = cc[0] >= 9
                except Exception:
                    hopper_capable = False
            if (
                not hopper_capable
                or not FLASH_ATTN_AVAILABLE
                or not hasattr(torch, "float8_e4m3fn")
            ):
                logger.info(
                    "FlashAttention3FP8: Hopper architecture (SM 9.0+), FlashAttention, or torch.float8_e4m3fn "
                    "not available; disabling FP8 and falling back to native SDPA."
                )
                self.use_fp8 = False

        # Static layer index
        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx: int = int(resolved_layer_id)

        # LoRA adaptation configuration
        self.lora_rank: int = getattr(
            config,
            "lora_rank",
            getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha: float = getattr(
            config,
            "lora_alpha",
            getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base: bool = self.lora_rank > 0

        # Decoder & nested encoder detection
        _mt = str(getattr(config, "model_type", "") or "").lower() if config is not None else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_dec = (
            (
                getattr(config, "is_decoder", False)
                or getattr(config, "is_causal", False)
                or any(_mt.startswith(p) for p in _decoder_prefixes)
            )
            if config is not None
            else False
        )
        if hasattr(config, "architectures") and any(
            "GPT" in str(a) or "CausalLM" in str(a)
            for a in (getattr(config, "architectures", []) or [])
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)

        is_nested = kwargs.get("is_nested_encoder")
        if is_nested is None and source_attn is not None:
            src_cls = source_attn.__class__.__name__
            if (
                src_cls
                in [
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                    "ElectraSelfAttention",
                    "DebertaSelfAttention",
                ]
                or (src_cls.endswith("SelfAttention") and "distil" not in src_cls.lower())
                or (hasattr(source_attn, "self") and hasattr(source_attn, "output"))
            ):
                is_nested = True
            elif (
                hasattr(source_attn, "out_lin")
                or hasattr(source_attn, "c_proj")
                or hasattr(source_attn, "o_proj")
            ):
                is_nested = False
        if is_nested is None:
            is_nested = False
        self.is_nested_encoder: bool = bool(is_nested)
        self.gradient_checkpointing: bool = (
            getattr(config, "gradient_checkpointing", False) if config is not None else False
        )

        # Multi-head projections equipped with LoRA adapters
        self.q_proj = LoRALinear(
            hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        self.rotary_emb = None
        self.comp_config: Any | None = comp_config

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    def copy_weights(self, src: Any) -> None:
        r"""Copy projection weights from pretrained attention module with Frobenius norm verification."""
        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 has_out_lin:
            is_nested = False
        elif (
            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())
        ):
            is_nested = True
        else:
            is_nested = False
        self.is_nested_encoder = bool(is_nested)

        if hasattr(src, "rotary_emb") and src.rotary_emb is not None:
            self.rotary_emb = src.rotary_emb

        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 + c_proj
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c_attn_mod = src.c_attn
            if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
                try:
                    c_attn_mod.merge()
                except Exception as e:
                    logger.warning(f"copy_weights c_attn merge warning: {e}")
            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()
                cls_c = getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__
                base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    (cls_c == "Conv1D")
                    or (base_cls_c == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                    or (
                        c_w.ndim == 2
                        and c_w.shape[0] < c_w.shape[1]
                        and "Linear" not in cls_c
                        and "Linear" not in base_cls_c
                    )
                ):
                    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()
                    cls_o = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__
                    base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__
                    if (
                        (cls_o == "Conv1D")
                        or (base_cls_o == "Conv1D")
                        or hasattr(c_proj_mod, "nf")
                        or hasattr(base_o, "nf")
                        or (
                            o_w.ndim == 2
                            and o_w.shape[0] < o_w.shape[1]
                            and "Linear" not in cls_o
                            and "Linear" not in base_cls_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. RoBERTa / BERT nested self-attention
        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. Standard linear projections (DistilBERT, Qwen, LLaMA)
        else:
            proj_candidates = {
                "q": ["q_lin", "q_proj", "query", "W_q"],
                "k": ["k_lin", "k_proj", "key", "W_k"],
                "v": ["v_lin", "v_proj", "value", "W_v"],
                "o": ["out_lin", "out_proj", "dense", "o_proj", "proj", "output"],
            }
            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

        # Fallback via UniversalAdapter
        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(f"UniversalAdapter fallback warning: {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"FlashAttention3FP8 copy_weights: missing mandatory projections {missing_after} "
                f"from source module {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

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

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        H, Dh = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        is_cross_attention = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross_attention else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Projections
        q = self.q_proj(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
        k = (
            self.k_proj(kv_input)
            .view(B, kv_len, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(kv_input)
            .view(B, kv_len, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. Rotary Position Embeddings (RoPE)
        if position_embeddings is not None and not is_cross_attention:
            if isinstance(position_embeddings, tuple):
                cos, sin = position_embeddings
            else:
                cos, sin = position_embeddings, None
            if sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        elif self.rotary_emb is not None and not is_cross_attention:
            if position_ids is not None:
                cos, sin = self.rotary_emb(v, seq_len=kv_len)
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)

        # 3. Grouped Query Attention (GQA) head expansion
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        causal_flag = bool(self.is_decoder and not is_cross_attention and kv_len == L)

        # 4. FlashAttention-3 FP8 hardware acceleration path
        # Requires Hopper architecture (SM 9.0+), flash_attn, and CUDA tensor
        out = None
        if FLASH_ATTN_AVAILABLE and self.use_fp8 and q.is_cuda and hasattr(torch, "float8_e4m3fn"):
            # Check Hopper compute capability
            is_hopper = False
            if hasattr(torch.cuda, "get_device_capability"):
                try:
                    is_hopper = torch.cuda.get_device_capability(device)[0] >= 9
                except Exception:
                    is_hopper = False

            if is_hopper:
                try:
                    # FlashAttention expects shape (B, L, H, Dh)
                    q_perm = q.transpose(1, 2).contiguous()
                    k_perm = k.transpose(1, 2).contiguous()
                    v_perm = v.transpose(1, 2).contiguous()

                    # Dynamic per-tensor scaling to FP8 max representable magnitude (448.0 for float8_e4m3fn)
                    fp8_max = 448.0
                    q_scale = (q_perm.abs().amax() / fp8_max).clamp(min=1e-6)
                    k_scale = (k_perm.abs().amax() / fp8_max).clamp(min=1e-6)
                    v_scale = (v_perm.abs().amax() / fp8_max).clamp(min=1e-6)

                    # Scale and cast to FP8
                    q_fa = (q_perm / q_scale).to(torch.float8_e4m3fn)
                    k_fa = (k_perm / k_scale).to(torch.float8_e4m3fn)
                    v_fa = (v_perm / v_scale).to(torch.float8_e4m3fn)

                    # Execute FlashAttention kernel
                    out_fp8 = flash_attn_func(q_fa, k_fa, v_fa, causal=causal_flag)

                    # De-scale back: output is mathematically scaled by v_scale
                    out_tensor = (out_fp8.float() * v_scale).to(dtype)

                    # Apply head_mask if provided
                    if head_mask is not None:
                        # out_tensor is (B, L, H, Dh) -> transpose to (B, H, L, Dh) for head masking
                        out_tensor_h = out_tensor.transpose(1, 2)
                        if head_mask.dim() == 1:
                            hm = head_mask.view(1, -1, 1, 1)
                        elif head_mask.dim() == 2:
                            hm = head_mask.view(head_mask.shape[0], head_mask.shape[1], 1, 1)
                        elif head_mask.dim() == 3:
                            hm = head_mask.unsqueeze(-1)
                        else:
                            hm = head_mask
                        out_tensor_h = out_tensor_h * hm
                        out = (
                            out_tensor_h.transpose(1, 2).contiguous().view(B, L, self.all_head_size)
                        )
                    else:
                        out = out_tensor.contiguous().view(B, L, self.all_head_size)
                except Exception as e:
                    logger.warning(
                        f"FlashAttention-3 FP8 execution failed: {e}; falling back to native SDPA."
                    )
                    out = None

        # 5. SDPA / PyTorch native fallback path
        if out is None:
            eff_mask = encoder_attention_mask if is_cross_attention else attention_mask
            attn_mask = None
            if eff_mask is not None:
                if eff_mask.dim() == 2:
                    min_val = torch.finfo(q.dtype).min if q.dtype.is_floating_point else -1e4
                    attn_mask = (1.0 - eff_mask.unsqueeze(1).unsqueeze(2).float()) * min_val
                elif eff_mask.dim() == 3:
                    attn_mask = eff_mask.unsqueeze(1).float()
                elif eff_mask.dim() == 4:
                    attn_mask = eff_mask.float()

            if causal_flag and attn_mask is None:
                out_sdpa = F.scaled_dot_product_attention(
                    q,
                    k,
                    v,
                    attn_mask=None,
                    is_causal=True,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
            else:
                if causal_flag:
                    causal_m = torch.triu(
                        torch.full(
                            (L, kv_len),
                            float("-inf"),
                            device=device,
                            dtype=torch.float32,
                        ),
                        diagonal=1,
                    )
                    attn_mask = (
                        (attn_mask + causal_m.unsqueeze(0).unsqueeze(0))
                        if attn_mask is not None
                        else causal_m.unsqueeze(0).unsqueeze(0)
                    )

                out_sdpa = F.scaled_dot_product_attention(
                    q,
                    k,
                    v,
                    attn_mask=attn_mask,
                    is_causal=False,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )

            # Head mask handling
            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.shape[0], head_mask.shape[1], 1, 1)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(-1)
                out_sdpa = out_sdpa * head_mask

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

        # 6. Output projection with nested encoder bypass
        if self.is_nested_encoder:
            final_out = out
        else:
            final_out = self.out_proj(out)

        return final_out, None

FlashAttention-3 FP8 Hardware-Accelerated Multi-Head Attention.

Accelerates attention using FP8 (e4m3fn) precision with FlashAttention-3 kernels on Hopper (SM 9.0+) GPUs. Automatically falls back to native PyTorch SDPA on pre-Hopper hardware or CPU.

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.

Instance variables

prop W_k : LoRALinear

Expand source code

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

prop W_q : LoRALinear

Expand source code

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

prop W_v : LoRALinear

Expand source code

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

prop c_proj : LoRALinear

Expand source code

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

prop dense : LoRALinear

Expand source code

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

prop k_lin : LoRALinear

Expand source code

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

prop key : LoRALinear

Expand source code

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

prop o_proj : LoRALinear

Expand source code

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

prop out_lin : LoRALinear

Expand source code

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

prop q_lin : LoRALinear

Expand source code

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

prop query : LoRALinear

Expand source code

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

prop v_lin : LoRALinear

Expand source code

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

prop value : LoRALinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    r"""Copy projection weights from pretrained attention module with Frobenius norm verification."""
    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 has_out_lin:
        is_nested = False
    elif (
        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())
    ):
        is_nested = True
    else:
        is_nested = False
    self.is_nested_encoder = bool(is_nested)

    if hasattr(src, "rotary_emb") and src.rotary_emb is not None:
        self.rotary_emb = src.rotary_emb

    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 + c_proj
    if hasattr(src, "c_attn") and src.c_attn is not None:
        c_attn_mod = src.c_attn
        if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False):
            try:
                c_attn_mod.merge()
            except Exception as e:
                logger.warning(f"copy_weights c_attn merge warning: {e}")
        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()
            cls_c = getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__
            base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__
            if (
                (cls_c == "Conv1D")
                or (base_cls_c == "Conv1D")
                or hasattr(c_attn_mod, "nf")
                or hasattr(base_c, "nf")
                or (
                    c_w.ndim == 2
                    and c_w.shape[0] < c_w.shape[1]
                    and "Linear" not in cls_c
                    and "Linear" not in base_cls_c
                )
            ):
                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()
                cls_o = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__
                base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    (cls_o == "Conv1D")
                    or (base_cls_o == "Conv1D")
                    or hasattr(c_proj_mod, "nf")
                    or hasattr(base_o, "nf")
                    or (
                        o_w.ndim == 2
                        and o_w.shape[0] < o_w.shape[1]
                        and "Linear" not in cls_o
                        and "Linear" not in base_cls_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. RoBERTa / BERT nested self-attention
    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. Standard linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q": ["q_lin", "q_proj", "query", "W_q"],
            "k": ["k_lin", "k_proj", "key", "W_k"],
            "v": ["v_lin", "v_proj", "value", "W_v"],
            "o": ["out_lin", "out_proj", "dense", "o_proj", "proj", "output"],
        }
        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

    # Fallback via UniversalAdapter
    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(f"UniversalAdapter fallback warning: {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"FlashAttention3FP8 copy_weights: missing mandatory projections {missing_after} "
            f"from source module {src.__class__.__name__}."
        )

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device, dtype=target_mod.weight.dtype
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

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

            if self.freeze_base:
                target_mod.freeze_base_weights()

Copy projection weights from pretrained attention module with Frobenius norm verification.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class FullAttention (config: Any,
rank: int | None = None,
lora_rank: int = 16,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "full", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "full_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class FullAttention(BaseAttention):
    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,
        rank: int | None = None,
        lora_rank: int = 16,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        self.adapter = ModelAdapter(config)
        self.hidden_size: int = self.adapter.hidden_size or kwargs.get(
            "hidden_size", kwargs.get("d_model", 768)
        )
        self.num_heads: int = self.adapter.num_heads or kwargs.get(
            "num_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads: int = self.adapter.num_key_value_heads or kwargs.get(
            "num_kv_heads", kwargs.get("num_key_value_heads", self.num_heads)
        )
        self.head_dim: int = self.adapter.head_dim or (self.hidden_size // max(1, self.num_heads))
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.is_decoder: bool = getattr(config, "is_decoder", False)
        self.is_nested_encoder: bool = False

        self.lora_rank = getattr(config, "lora_rank", lora_rank)

        self.query = LoRALinear(
            self.hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.key = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.value = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                ),
            )
            if config
            else 0.1
        )
        self.dropout: nn.Dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

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

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

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

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

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

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

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

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

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

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

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

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # 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
        ):
            weight = src.c_attn.weight.data
            bias = getattr(src.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                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:
            w_o = src.c_proj.weight.data
            if (
                w_o.dim() == 2
                and w_o.shape[0] != self.hidden_size
                and w_o.shape[1] == self.hidden_size
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(src.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            (
                hasattr(src, "self")
                and hasattr(src.self, "query")
                and hasattr(src, "output")
                and hasattr(src.output, "dense")
            )
            or (
                cls_name in ("BertSelfAttention", "RobertaSelfAttention")
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        # Check required weights
        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(
                "FullAttention copy_weights warning. Proceeding with initialized weights."
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        "FullAttention extracted weight for %s has near-zero Frobenius norm (%.6f)",
                        proj_key,
                        norm.item(),
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    "FullAttention extracted weight for output projection has near-zero Frobenius norm (%.6f)",
                    norm_o.item(),
                )

        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

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

                target.freeze_base_weights()

        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

        Q = self.query(hidden_states).view(B, N_q, H, d).transpose(1, 2).contiguous()
        K = self.key(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V = self.value(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # Apply RoPE
        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"))

        # 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
        scores = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale

        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)

        # Explicit Causal Masking for decoders
        is_decoder = (
            self.is_decoder
            or getattr(self.config, "is_decoder", False)
            or getattr(self.config, "model_type", "")
            in ("gpt2", "qwen", "qwen2", "qwen3_5", "llama")
        )
        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)

        attn = torch.nan_to_num(F.softmax(scores.float(), dim=-1), nan=0.0).to(dtype)
        attn = self.dropout(attn)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn = attn * head_mask

        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)

        return output, (attn if output_attentions else None)

Every 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_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 : LoRALinear

Expand source code

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

prop dense : LoRALinear

Expand source code

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

prop k_lin : LoRALinear

Expand source code

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

prop k_proj : LoRALinear

Expand source code

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

prop o_proj : LoRALinear

Expand source code

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

prop out_lin : LoRALinear

Expand source code

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

prop q_lin : LoRALinear

Expand source code

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

prop q_proj : LoRALinear

Expand source code

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

prop v_lin : LoRALinear

Expand source code

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

prop v_proj : LoRALinear

Expand source code

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

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class FusedSVDKernelAttention (config: Any | None = None,
rank: int | None = None,
old_attn: nn.Module | None = None,
use_fused_kernel: bool = True,
use_rms_norm: bool = False,
dropout: float = 0.1,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "fused_svd_kernel",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "fused_svd_kernel_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class FusedSVDKernelAttention(BaseAttention):
    r"""Fused SVD-Kernel Linear Attention Module.
    Production-hardened across DistilBERT, RoBERTa/BERT nested, GPT-2 Conv1D, and Qwen/LLaMA.
    """

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

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        old_attn: nn.Module | None = None,
        use_fused_kernel: bool = True,
        use_rms_norm: bool = False,
        dropout: float = 0.1,
        **kwargs: Any,
    ) -> None:
        h_dim = int(
            kwargs.get("hidden_size")
            or _cfg_get(config, "hidden_size", None)
            or _cfg_get(config, "dim", None)
            or _cfg_get(config, "d_model", None)
            or _cfg_get(config, "n_embd", None)
            or getattr(
                old_attn if old_attn is not None else kwargs.get("pretrained_attn"),
                "embed_dim",
                None,
            )
            or getattr(
                old_attn if old_attn is not None else kwargs.get("pretrained_attn"),
                "hidden_size",
                None,
            )
            or 768
        )
        self.hidden_size: int = h_dim
        actual_rank = rank or _cfg_get(
            config, "rank", _cfg_get(config, "svd_rank", self.hidden_size)
        )
        super().__init__(config, rank=actual_rank)

        src_attn = old_attn if old_attn is not None else kwargs.get("pretrained_attn")
        self.num_heads: int = int(
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or _cfg_get(config, "num_attention_heads", None)
            or _cfg_get(config, "n_heads", None)
            or _cfg_get(config, "num_heads", None)
            or _cfg_get(config, "n_head", None)
            or getattr(src_attn, "num_attention_heads", None)
            or getattr(src_attn, "n_heads", None)
            or getattr(src_attn, "num_heads", None)
            or (self.hidden_size // 64 if self.hidden_size % 64 == 0 else 12)
        )
        self.head_dim: int = int(
            kwargs.get("head_dim")
            or _cfg_get(
                config,
                "head_dim",
                self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
            )
        )
        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.q_head_size: int = self.num_heads * self.head_dim
        self.all_head_size: int = self.q_head_size

        self.rank: int = actual_rank
        self.use_fused_kernel: bool = use_fused_kernel
        self.use_rms_norm: bool = use_rms_norm

        is_dec = (
            bool(_cfg_get(config, "is_decoder", False))
            or any(
                str(_cfg_get(config, "model_type", "")).lower().startswith(p)
                for p in (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "gemma",
                    "phi",
                    "opt",
                    "bloom",
                    "falcon",
                    "olmo",
                    "mpt",
                    "stablelm",
                )
            )
            or (src_attn is not None and hasattr(src_attn, "c_attn"))
            or bool(kwargs.get("is_causal", False))
            or bool(kwargs.get("is_decoder", False))
        )
        self.is_decoder: bool = is_dec
        self.is_causal: bool = is_dec
        self.is_nested_encoder: bool = False

        self.dropout_rate: float = float(
            _cfg_get(
                config,
                "attention_probs_dropout_prob",
                _cfg_get(config, "dropout", dropout),
            )
        )
        self.dropout = nn.Dropout(self.dropout_rate)

        self.q_proj = FusedSVDLinear(
            self.hidden_size,
            self.q_head_size,
            rank=self.rank,
            bias=True,
            use_rms_norm=self.use_rms_norm,
        )
        self.k_proj = FusedSVDLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            use_rms_norm=self.use_rms_norm,
        )
        self.v_proj = FusedSVDLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            use_rms_norm=self.use_rms_norm,
        )
        self.out_proj = FusedSVDLinear(
            self.q_head_size,
            self.hidden_size,
            rank=self.rank,
            bias=True,
            use_rms_norm=self.use_rms_norm,
        )

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

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

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

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

    def copy_weights(self, old_attn: Any) -> None:
        """Multi-architecture weight copying supporting GPT-2 Conv1D, RoBERTa nested, and flat projections."""
        if old_attn is None:
            return

        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_out = (
                hasattr(old_attn, "out_lin")
                or hasattr(old_attn, "dense")
                or hasattr(old_attn, "out_proj")
                or hasattr(old_attn, "c_proj")
                or hasattr(old_attn, "o_proj")
            )

            if (
                hasattr(old_attn, "out_lin")
                or hasattr(old_attn, "c_proj")
                or hasattr(old_attn, "o_proj")
            ):
                self.is_nested_encoder = False
            elif (
                (
                    hasattr(old_attn, "self")
                    and hasattr(old_attn.self, "query")
                    and hasattr(old_attn, "output")
                    and hasattr(old_attn.output, "dense")
                )
                or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention")
                or (hasattr(old_attn, "query") and not has_out)
            ):
                self.is_nested_encoder = True
            else:
                self.is_nested_encoder = False

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                if (
                    hasattr(old_attn, "output")
                    and hasattr(old_attn.output, "dense")
                    and hasattr(old_attn.output.dense, "weight")
                ):
                    w_o = old_attn.output.dense.weight.data
                    b_o = getattr(old_attn.output.dense, "bias", None)
                    b_o = b_o.data if b_o is not None else None
            src = old_attn
            while hasattr(src, "self") and not any(
                hasattr(src, attr)
                for attr in ("query", "key", "value", "q_proj", "k_proj", "v_proj")
            ):
                self.is_nested_encoder = True
                src = src.self

            # 1. GPT-2 Conv1D fused c_attn & c_proj
            if hasattr(src, "c_attn") and src.c_attn is not None:
                self.is_decoder = True
                self.is_causal = True
                mod_c = src.c_attn
                base_c = getattr(mod_c, "base_layer", mod_c)
                w = getattr(base_c, "weight", getattr(mod_c, "weight", None))
                if w is not None:
                    w = w.data.clone()
                    cls_c = getattr(mod_c, "__class__", type(mod_c)).__name__
                    base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__
                    if (
                        "Conv1D" in cls_c
                        or "Conv1D" in base_cls_c
                        or hasattr(mod_c, "nf")
                        or hasattr(base_c, "nf")
                        or (w.ndim == 2 and w.shape[0] < w.shape[1])
                    ):
                        w = w.t().contiguous()
                    q_dim = self.q_head_size
                    k_dim = self.kv_head_size
                    v_dim = self.kv_head_size
                    if w.shape[0] == q_dim + k_dim + v_dim:
                        w_q = w[:q_dim, :]
                        w_k = w[q_dim : q_dim + k_dim, :]
                        w_v = w[q_dim + k_dim :, :]
                    else:
                        split_size = w.shape[0] // 3
                        w_q = w[:split_size, :]
                        w_k = w[split_size : 2 * split_size, :]
                        w_v = w[2 * split_size :, :]
                    b = getattr(base_c, "bias", getattr(mod_c, "bias", None))
                    if b is not None:
                        b_data = b.data.clone()
                        if b_data.shape[0] == q_dim + k_dim + v_dim:
                            b_q = b_data[:q_dim]
                            b_k = b_data[q_dim : q_dim + k_dim]
                            b_v = b_data[q_dim + k_dim :]
                        else:
                            b_q = b_data[:split_size]
                            b_k = b_data[split_size : 2 * split_size]
                            b_v = b_data[2 * split_size :]

                if hasattr(src, "c_proj") and src.c_proj is not None:
                    mod_o = src.c_proj
                    base_o = getattr(mod_o, "base_layer", mod_o)
                    out_w = getattr(base_o, "weight", getattr(mod_o, "weight", None))
                    if out_w is not None:
                        out_w = out_w.data.clone()
                        cls_o = getattr(mod_o, "__class__", type(mod_o)).__name__
                        base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__
                        if (
                            "Conv1D" in cls_o
                            or "Conv1D" in base_cls_o
                            or hasattr(mod_o, "nf")
                            or hasattr(base_o, "nf")
                            or out_w.ndim == 2
                            and out_w.shape[0] < out_w.shape[1]
                        ):
                            out_w = out_w.t().contiguous()
                        w_o = out_w
                        out_b = getattr(base_o, "bias", getattr(mod_o, "bias", None))
                        b_o = out_b.data.clone() if out_b is not None else None

            # 2. DistilBERT flat projections (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.data.clone()
                b_q = (
                    src.q_lin.bias.data.clone()
                    if getattr(src.q_lin, "bias", None) is not None
                    else None
                )
                w_k = src.k_lin.weight.data.clone()
                b_k = (
                    src.k_lin.bias.data.clone()
                    if getattr(src.k_lin, "bias", None) is not None
                    else None
                )
                w_v = src.v_lin.weight.data.clone()
                b_v = (
                    src.v_lin.bias.data.clone()
                    if getattr(src.v_lin, "bias", None) is not None
                    else None
                )
                if (
                    hasattr(src, "out_lin")
                    and getattr(src, "out_lin", None) is not None
                    and hasattr(src.out_lin, "weight")
                ):
                    w_o = src.out_lin.weight.data.clone()
                    b_o = (
                        src.out_lin.bias.data.clone()
                        if getattr(src.out_lin, "bias", None) is not None
                        else None
                    )

            # 3. Flat projection splitting (single in_proj_weight or qkv_proj / Wqkv / qkv)
            elif (
                (
                    hasattr(src, "in_proj_weight")
                    and getattr(src, "in_proj_weight", None) is not None
                )
                or (
                    hasattr(src, "qkv_proj")
                    and getattr(src, "qkv_proj", None) is not None
                    and hasattr(src.qkv_proj, "weight")
                )
                or (
                    hasattr(src, "qkv")
                    and getattr(src, "qkv", None) is not None
                    and hasattr(src.qkv, "weight")
                )
                or (
                    hasattr(src, "Wqkv")
                    and getattr(src, "Wqkv", None) is not None
                    and hasattr(src.Wqkv, "weight")
                )
            ):
                qkv_mod = getattr(src, "qkv_proj", getattr(src, "qkv", getattr(src, "Wqkv", None)))
                if qkv_mod is not None:
                    in_w = qkv_mod.weight.data.clone()
                    in_b = (
                        qkv_mod.bias.data.clone()
                        if getattr(qkv_mod, "bias", None) is not None
                        else None
                    )
                else:
                    in_w = src.in_proj_weight.data.clone()
                    in_b = getattr(src, "in_proj_bias", None)
                    in_b = in_b.data.clone() if in_b is not None else None

                if in_w.ndim == 2 and in_w.shape[0] < in_w.shape[1]:
                    in_w = in_w.t().contiguous()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size
                if in_w.shape[0] == q_dim + k_dim + v_dim:
                    w_q = in_w[:q_dim, :]
                    w_k = in_w[q_dim : q_dim + k_dim, :]
                    w_v = in_w[q_dim + k_dim :, :]
                else:
                    split_sz = in_w.shape[0] // 3
                    w_q = in_w[:split_sz, :]
                    w_k = in_w[split_sz : 2 * split_sz, :]
                    w_v = in_w[2 * split_sz :, :]

                if in_b is not None:
                    if in_b.shape[0] == q_dim + k_dim + v_dim:
                        b_q = in_b[:q_dim]
                        b_k = in_b[q_dim : q_dim + k_dim]
                        b_v = in_b[q_dim + k_dim :]
                    else:
                        split_sz = in_b.shape[0] // 3
                        b_q = in_b[:split_sz]
                        b_k = in_b[split_sz : 2 * split_sz]
                        b_v = in_b[2 * split_sz :]

                for out_name in ("out_proj", "o_proj", "dense", "out_lin"):
                    if (
                        hasattr(src, out_name)
                        and getattr(src, out_name, None) is not None
                        and hasattr(getattr(src, out_name), "weight")
                    ):
                        m_out = getattr(src, out_name)
                        w_o = m_out.weight.data.clone()
                        b_o = (
                            m_out.bias.data.clone()
                            if getattr(m_out, "bias", None) is not None
                            else None
                        )
                        break

            # 4. RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.data.clone()
                b_q = (
                    src.query.bias.data.clone()
                    if getattr(src.query, "bias", None) is not None
                    else None
                )
                w_k = src.key.weight.data.clone()
                b_k = (
                    src.key.bias.data.clone()
                    if getattr(src.key, "bias", None) is not None
                    else None
                )
                w_v = src.value.weight.data.clone()
                b_v = (
                    src.value.bias.data.clone()
                    if getattr(src.value, "bias", None) is not None
                    else None
                )
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if (
                            hasattr(src, attr)
                            and getattr(src, attr) is not None
                            and hasattr(getattr(src, attr), "weight")
                        ):
                            m = getattr(src, attr)
                            w_o = m.weight.data.clone()
                            b_o = (
                                m.bias.data.clone()
                                if getattr(m, "bias", None) is not None
                                else None
                            )
                            break

            # 5. Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.data.clone()
                b_q = (
                    src.q_proj.bias.data.clone()
                    if getattr(src.q_proj, "bias", None) is not None
                    else None
                )
                w_k = src.k_proj.weight.data.clone()
                b_k = (
                    src.k_proj.bias.data.clone()
                    if getattr(src.k_proj, "bias", None) is not None
                    else None
                )
                w_v = src.v_proj.weight.data.clone()
                b_v = (
                    src.v_proj.bias.data.clone()
                    if getattr(src.v_proj, "bias", None) is not None
                    else None
                )
                if (
                    hasattr(src, "o_proj")
                    and getattr(src, "o_proj", None) is not None
                    and hasattr(src.o_proj, "weight")
                ):
                    w_o = src.o_proj.weight.data.clone()
                    b_o = (
                        src.o_proj.bias.data.clone()
                        if getattr(src.o_proj, "bias", None) is not None
                        else None
                    )

            # 6. Fallback: UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    from neuralops.core.adapters import UniversalAdapter

                    adapter = UniversalAdapter(old_attn)
                    extracted = adapter.extract_qkv_out()
                    if extracted:
                        w_q = extracted.get("q_weight")
                        b_q = extracted.get("q_bias")
                        w_k = extracted.get("k_weight")
                        b_k = extracted.get("k_bias")
                        w_v = extracted.get("v_weight")
                        b_v = extracted.get("v_bias")
                        if w_o is None:
                            w_o = extracted.get("out_weight")
                            b_o = extracted.get("out_bias")
                except Exception as ex:
                    logger.debug("UniversalAdapter fallback failed: %s", ex)

            missing = []
            if w_q is None:
                missing.append("q")
            if w_k is None:
                missing.append("k")
            if w_v is None:
                missing.append("v")
            if missing:
                logger.warning(
                    f"Failed to extract required projections {missing} from {src_cls_name}"
                )

            # Frobenius norm verification >= 1e-4 on copied weights
            for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
                if w is not None and torch.linalg.norm(w.float()).item() < 1e-4:
                    logger.warning(f"Copied weight {name} failed Frobenius norm check (< 1e-4)")
            if w_o is not None and not self.is_nested_encoder:
                if torch.linalg.norm(w_o.float()).item() < 1e-4:
                    logger.warning("Copied weight w_o failed Frobenius norm check (< 1e-4)")

            # SVD decomposition into low-rank factors
            self.q_proj.decompose_from_weight(w_q, b_q)
            self.k_proj.decompose_from_weight(w_k, b_k)
            self.v_proj.decompose_from_weight(w_v, b_v)
            if w_o is not None:
                self.out_proj.decompose_from_weight(w_o, b_o)
            elif self.is_nested_encoder:
                with torch.no_grad():
                    U = self.out_proj.U.data
                    V = self.out_proj.V.data
                    U.zero_()
                    V.zero_()
                    d = min(U.shape)
                    U[:d, :d].fill_diagonal_(1.0)
                    d2 = min(V.shape)
                    V[:d2, :d2].fill_diagonal_(1.0)
                    if self.out_proj.bias is not None:
                        self.out_proj.bias.data.zero_()

    def forward(
        self,
        hidden_states: torch.Tensor | None = None,
        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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("Cross-attention is not supported in FusedSVDKernelAttention")

        x = hidden_states if hidden_states is not None else kwargs.get("query", kwargs.get("x"))
        if x is None and len(args) > 0:
            x = args[0]
        if attention_mask is None:
            attention_mask = kwargs.get("mask", kwargs.get("attn_mask"))

        B, L, _ = x.shape

        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k = self.k_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        v = self.v_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()

        if position_embeddings is not None:
            cos, sin = position_embeddings
            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)

        scale = 1.0 / math.sqrt(self.head_dim)
        # Memory-efficient: use SDPA to avoid OOM on large seqs

        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 = (
                getattr(self, "is_decoder", False) or getattr(self, "is_causal", False)
            ) and L > 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(x.dtype).transpose(1, 2).contiguous().view(B, L, self.all_head_size)
            )
            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None
        except Exception:
            pass

        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder or self.is_causal:
            causal_mask = torch.tril(torch.ones(L, L, device=x.device, dtype=torch.bool))
            scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4)

        if attention_mask is not None:
            if attention_mask.dim() == 2:
                valid_mask = get_token_valid_mask(attention_mask, B, L, x.device)
                if valid_mask is None:
                    valid_mask = torch.ones(B, L, device=x.device, dtype=torch.bool)
                scores = scores.masked_fill(~valid_mask.unsqueeze(1).unsqueeze(2), -1e4)
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask.clamp(min=-1e4)

        probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(v.dtype)
        probs = self.dropout(probs)

        if head_mask is not None:
            probs = probs * head_mask

        context = torch.matmul(probs, v).transpose(1, 2).contiguous().view(B, L, self.q_head_size)

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

        return out, (probs if output_attentions else None)

Fused SVD-Kernel Linear Attention Module. Production-hardened across DistilBERT, RoBERTa/BERT nested, GPT-2 Conv1D, and Qwen/LLaMA.

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.

Instance variables

prop key : FusedSVDLinear

Expand source code

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

prop query : FusedSVDLinear

Expand source code

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

prop value : FusedSVDLinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Multi-architecture weight copying supporting GPT-2 Conv1D, RoBERTa nested, and flat projections."""
    if old_attn is None:
        return

    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        has_out = (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "dense")
            or hasattr(old_attn, "out_proj")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        )

        if (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
        ):
            self.is_nested_encoder = False
        elif (
            (
                hasattr(old_attn, "self")
                and hasattr(old_attn.self, "query")
                and hasattr(old_attn, "output")
                and hasattr(old_attn.output, "dense")
            )
            or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (hasattr(old_attn, "query") and not has_out)
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            if (
                hasattr(old_attn, "output")
                and hasattr(old_attn.output, "dense")
                and hasattr(old_attn.output.dense, "weight")
            ):
                w_o = old_attn.output.dense.weight.data
                b_o = getattr(old_attn.output.dense, "bias", None)
                b_o = b_o.data if b_o is not None else None
        src = old_attn
        while hasattr(src, "self") and not any(
            hasattr(src, attr)
            for attr in ("query", "key", "value", "q_proj", "k_proj", "v_proj")
        ):
            self.is_nested_encoder = True
            src = src.self

        # 1. GPT-2 Conv1D fused c_attn & c_proj
        if hasattr(src, "c_attn") and src.c_attn is not None:
            self.is_decoder = True
            self.is_causal = True
            mod_c = src.c_attn
            base_c = getattr(mod_c, "base_layer", mod_c)
            w = getattr(base_c, "weight", getattr(mod_c, "weight", None))
            if w is not None:
                w = w.data.clone()
                cls_c = getattr(mod_c, "__class__", type(mod_c)).__name__
                base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    "Conv1D" in cls_c
                    or "Conv1D" in base_cls_c
                    or hasattr(mod_c, "nf")
                    or hasattr(base_c, "nf")
                    or (w.ndim == 2 and w.shape[0] < w.shape[1])
                ):
                    w = w.t().contiguous()
                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size
                if w.shape[0] == q_dim + k_dim + v_dim:
                    w_q = w[:q_dim, :]
                    w_k = w[q_dim : q_dim + k_dim, :]
                    w_v = w[q_dim + k_dim :, :]
                else:
                    split_size = w.shape[0] // 3
                    w_q = w[:split_size, :]
                    w_k = w[split_size : 2 * split_size, :]
                    w_v = w[2 * split_size :, :]
                b = getattr(base_c, "bias", getattr(mod_c, "bias", None))
                if b is not None:
                    b_data = b.data.clone()
                    if b_data.shape[0] == q_dim + k_dim + v_dim:
                        b_q = b_data[:q_dim]
                        b_k = b_data[q_dim : q_dim + k_dim]
                        b_v = b_data[q_dim + k_dim :]
                    else:
                        b_q = b_data[:split_size]
                        b_k = b_data[split_size : 2 * split_size]
                        b_v = b_data[2 * split_size :]

            if hasattr(src, "c_proj") and src.c_proj is not None:
                mod_o = src.c_proj
                base_o = getattr(mod_o, "base_layer", mod_o)
                out_w = getattr(base_o, "weight", getattr(mod_o, "weight", None))
                if out_w is not None:
                    out_w = out_w.data.clone()
                    cls_o = getattr(mod_o, "__class__", type(mod_o)).__name__
                    base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__
                    if (
                        "Conv1D" in cls_o
                        or "Conv1D" in base_cls_o
                        or hasattr(mod_o, "nf")
                        or hasattr(base_o, "nf")
                        or out_w.ndim == 2
                        and out_w.shape[0] < out_w.shape[1]
                    ):
                        out_w = out_w.t().contiguous()
                    w_o = out_w
                    out_b = getattr(base_o, "bias", getattr(mod_o, "bias", None))
                    b_o = out_b.data.clone() if out_b is not None else None

        # 2. DistilBERT flat projections (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.data.clone()
            b_q = (
                src.q_lin.bias.data.clone()
                if getattr(src.q_lin, "bias", None) is not None
                else None
            )
            w_k = src.k_lin.weight.data.clone()
            b_k = (
                src.k_lin.bias.data.clone()
                if getattr(src.k_lin, "bias", None) is not None
                else None
            )
            w_v = src.v_lin.weight.data.clone()
            b_v = (
                src.v_lin.bias.data.clone()
                if getattr(src.v_lin, "bias", None) is not None
                else None
            )
            if (
                hasattr(src, "out_lin")
                and getattr(src, "out_lin", None) is not None
                and hasattr(src.out_lin, "weight")
            ):
                w_o = src.out_lin.weight.data.clone()
                b_o = (
                    src.out_lin.bias.data.clone()
                    if getattr(src.out_lin, "bias", None) is not None
                    else None
                )

        # 3. Flat projection splitting (single in_proj_weight or qkv_proj / Wqkv / qkv)
        elif (
            (
                hasattr(src, "in_proj_weight")
                and getattr(src, "in_proj_weight", None) is not None
            )
            or (
                hasattr(src, "qkv_proj")
                and getattr(src, "qkv_proj", None) is not None
                and hasattr(src.qkv_proj, "weight")
            )
            or (
                hasattr(src, "qkv")
                and getattr(src, "qkv", None) is not None
                and hasattr(src.qkv, "weight")
            )
            or (
                hasattr(src, "Wqkv")
                and getattr(src, "Wqkv", None) is not None
                and hasattr(src.Wqkv, "weight")
            )
        ):
            qkv_mod = getattr(src, "qkv_proj", getattr(src, "qkv", getattr(src, "Wqkv", None)))
            if qkv_mod is not None:
                in_w = qkv_mod.weight.data.clone()
                in_b = (
                    qkv_mod.bias.data.clone()
                    if getattr(qkv_mod, "bias", None) is not None
                    else None
                )
            else:
                in_w = src.in_proj_weight.data.clone()
                in_b = getattr(src, "in_proj_bias", None)
                in_b = in_b.data.clone() if in_b is not None else None

            if in_w.ndim == 2 and in_w.shape[0] < in_w.shape[1]:
                in_w = in_w.t().contiguous()

            q_dim = self.q_head_size
            k_dim = self.kv_head_size
            v_dim = self.kv_head_size
            if in_w.shape[0] == q_dim + k_dim + v_dim:
                w_q = in_w[:q_dim, :]
                w_k = in_w[q_dim : q_dim + k_dim, :]
                w_v = in_w[q_dim + k_dim :, :]
            else:
                split_sz = in_w.shape[0] // 3
                w_q = in_w[:split_sz, :]
                w_k = in_w[split_sz : 2 * split_sz, :]
                w_v = in_w[2 * split_sz :, :]

            if in_b is not None:
                if in_b.shape[0] == q_dim + k_dim + v_dim:
                    b_q = in_b[:q_dim]
                    b_k = in_b[q_dim : q_dim + k_dim]
                    b_v = in_b[q_dim + k_dim :]
                else:
                    split_sz = in_b.shape[0] // 3
                    b_q = in_b[:split_sz]
                    b_k = in_b[split_sz : 2 * split_sz]
                    b_v = in_b[2 * split_sz :]

            for out_name in ("out_proj", "o_proj", "dense", "out_lin"):
                if (
                    hasattr(src, out_name)
                    and getattr(src, out_name, None) is not None
                    and hasattr(getattr(src, out_name), "weight")
                ):
                    m_out = getattr(src, out_name)
                    w_o = m_out.weight.data.clone()
                    b_o = (
                        m_out.bias.data.clone()
                        if getattr(m_out, "bias", None) is not None
                        else None
                    )
                    break

        # 4. RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.data.clone()
            b_q = (
                src.query.bias.data.clone()
                if getattr(src.query, "bias", None) is not None
                else None
            )
            w_k = src.key.weight.data.clone()
            b_k = (
                src.key.bias.data.clone()
                if getattr(src.key, "bias", None) is not None
                else None
            )
            w_v = src.value.weight.data.clone()
            b_v = (
                src.value.bias.data.clone()
                if getattr(src.value, "bias", None) is not None
                else None
            )
            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if (
                        hasattr(src, attr)
                        and getattr(src, attr) is not None
                        and hasattr(getattr(src, attr), "weight")
                    ):
                        m = getattr(src, attr)
                        w_o = m.weight.data.clone()
                        b_o = (
                            m.bias.data.clone()
                            if getattr(m, "bias", None) is not None
                            else None
                        )
                        break

        # 5. Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.data.clone()
            b_q = (
                src.q_proj.bias.data.clone()
                if getattr(src.q_proj, "bias", None) is not None
                else None
            )
            w_k = src.k_proj.weight.data.clone()
            b_k = (
                src.k_proj.bias.data.clone()
                if getattr(src.k_proj, "bias", None) is not None
                else None
            )
            w_v = src.v_proj.weight.data.clone()
            b_v = (
                src.v_proj.bias.data.clone()
                if getattr(src.v_proj, "bias", None) is not None
                else None
            )
            if (
                hasattr(src, "o_proj")
                and getattr(src, "o_proj", None) is not None
                and hasattr(src.o_proj, "weight")
            ):
                w_o = src.o_proj.weight.data.clone()
                b_o = (
                    src.o_proj.bias.data.clone()
                    if getattr(src.o_proj, "bias", None) is not None
                    else None
                )

        # 6. Fallback: UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                from neuralops.core.adapters import UniversalAdapter

                adapter = UniversalAdapter(old_attn)
                extracted = adapter.extract_qkv_out()
                if extracted:
                    w_q = extracted.get("q_weight")
                    b_q = extracted.get("q_bias")
                    w_k = extracted.get("k_weight")
                    b_k = extracted.get("k_bias")
                    w_v = extracted.get("v_weight")
                    b_v = extracted.get("v_bias")
                    if w_o is None:
                        w_o = extracted.get("out_weight")
                        b_o = extracted.get("out_bias")
            except Exception as ex:
                logger.debug("UniversalAdapter fallback failed: %s", ex)

        missing = []
        if w_q is None:
            missing.append("q")
        if w_k is None:
            missing.append("k")
        if w_v is None:
            missing.append("v")
        if missing:
            logger.warning(
                f"Failed to extract required projections {missing} from {src_cls_name}"
            )

        # Frobenius norm verification >= 1e-4 on copied weights
        for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
            if w is not None and torch.linalg.norm(w.float()).item() < 1e-4:
                logger.warning(f"Copied weight {name} failed Frobenius norm check (< 1e-4)")
        if w_o is not None and not self.is_nested_encoder:
            if torch.linalg.norm(w_o.float()).item() < 1e-4:
                logger.warning("Copied weight w_o failed Frobenius norm check (< 1e-4)")

        # SVD decomposition into low-rank factors
        self.q_proj.decompose_from_weight(w_q, b_q)
        self.k_proj.decompose_from_weight(w_k, b_k)
        self.v_proj.decompose_from_weight(w_v, b_v)
        if w_o is not None:
            self.out_proj.decompose_from_weight(w_o, b_o)
        elif self.is_nested_encoder:
            with torch.no_grad():
                U = self.out_proj.U.data
                V = self.out_proj.V.data
                U.zero_()
                V.zero_()
                d = min(U.shape)
                U[:d, :d].fill_diagonal_(1.0)
                d2 = min(V.shape)
                V[:d2, :d2].fill_diagonal_(1.0)
                if self.out_proj.bias is not None:
                    self.out_proj.bias.data.zero_()

Multi-architecture weight copying supporting GPT-2 Conv1D, RoBERTa nested, and flat projections.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class FusedSVDLinear (in_features: int,
out_features: int,
rank: int = 32,
bias: bool = True,
use_rms_norm: bool = False,
eps: float = 1e-06,
dtype: torch.dtype = torch.float32,
device: torch.device | None = None)

Expand source code

class FusedSVDLinear(nn.Module):
    """
    RESEARCH SOTA: Fused SVD-Linear Module.
    Combines low-rank projection $Z = x V^T$, in-SRAM RMSNorm, and up-projection $Y = Z U^T + b$
    into a single kernel invocation, eliminating HBM roundtrips.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 32,
        bias: bool = True,
        use_rms_norm: bool = False,
        eps: float = 1e-6,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.eps = eps
        self.use_rms_norm = use_rms_norm

        # Low-rank factorized parameters
        self.V = nn.Parameter(torch.empty(rank, in_features, dtype=dtype, device=device))
        self.U = nn.Parameter(torch.empty(out_features, rank, dtype=dtype, device=device))

        if use_rms_norm:
            self.norm_weight = nn.Parameter(torch.ones(rank, dtype=dtype, device=device))
        else:
            self.register_parameter("norm_weight", None)

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        # Symmetrical square-root initialization
        nn.init.orthogonal_(self.U)
        nn.init.orthogonal_(self.V)
        if self.bias is not None:
            nn.init.zeros_(self.bias)
        if self.norm_weight is not None:
            nn.init.ones_(self.norm_weight)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        return torch.matmul(self.U, self.V)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # RESEARCH FIX: Dispatches fused kernel executing down-projection, RMSNorm, and up-projection
        return fused_svd_linear(
            x=x,
            V=self.V,
            U=self.U,
            norm_weight=self.norm_weight,
            bias=self.bias,
            eps=self.eps,
            use_rms_norm=self.use_rms_norm,
        )

    def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
        """Set precomputed U and V low-rank factors."""
        with torch.no_grad():
            self.U.copy_(U_factor)
            self.V.copy_(V_factor)

    def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
        """Decompose dense weight matrix into symmetric SVD factors."""
        with torch.no_grad():
            W = weight.float()
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            r = min(self.rank, len(S))
            sqrt_s = torch.sqrt(S[:r])
            self.U.data.zero_()
            self.V.data.zero_()
            self.U.data[:, :r] = (U[:, :r] * sqrt_s).to(self.U.dtype)
            self.V.data[:r, :] = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(self.V.dtype)
            if bias is not None and self.bias is not None:
                self.bias.data.copy_(bias)

    def get_orthogonality_loss(self) -> torch.Tensor:
        """
        # RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse.
        """
        u_gram = torch.matmul(self.U.T, self.U)
        v_gram = torch.matmul(self.V, self.V.T)
        eye_r = torch.eye(self.rank, device=self.U.device, dtype=self.U.dtype)
        loss_u = torch.norm(u_gram - eye_r, p="fro") ** 2
        loss_v = torch.norm(v_gram - eye_r, p="fro") ** 2
        return (loss_u + loss_v) / (self.rank * self.rank)

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"rank={self.rank}, rms_norm={self.use_rms_norm}, bias={self.bias is not None}"
        )

RESEARCH SOTA: Fused SVD-Linear Module. Combines low-rank projection $Z = x V^T$, in-SRAM RMSNorm, and up-projection $Y = Z U^T + b$ into a single kernel invocation, eliminating HBM roundtrips.

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

Ancestors

• torch.nn.modules.module.Module

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    return torch.matmul(self.U, self.V)

Reconstructed effective weight matrix for compatibility.

Methods

def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) ‑> None

Expand source code

def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
    """Decompose dense weight matrix into symmetric SVD factors."""
    with torch.no_grad():
        W = weight.float()
        U, S, Vt = torch.linalg.svd(W, full_matrices=False)
        r = min(self.rank, len(S))
        sqrt_s = torch.sqrt(S[:r])
        self.U.data.zero_()
        self.V.data.zero_()
        self.U.data[:, :r] = (U[:, :r] * sqrt_s).to(self.U.dtype)
        self.V.data[:r, :] = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(self.V.dtype)
        if bias is not None and self.bias is not None:
            self.bias.data.copy_(bias)

Decompose dense weight matrix into symmetric SVD factors.

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, out_features={self.out_features}, "
        f"rank={self.rank}, rms_norm={self.use_rms_norm}, bias={self.bias is not None}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # RESEARCH FIX: Dispatches fused kernel executing down-projection, RMSNorm, and up-projection
    return fused_svd_linear(
        x=x,
        V=self.V,
        U=self.U,
        norm_weight=self.norm_weight,
        bias=self.bias,
        eps=self.eps,
        use_rms_norm=self.use_rms_norm,
    )

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def get_orthogonality_loss(self) ‑> torch.Tensor

Expand source code

def get_orthogonality_loss(self) -> torch.Tensor:
    """
    # RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse.
    """
    u_gram = torch.matmul(self.U.T, self.U)
    v_gram = torch.matmul(self.V, self.V.T)
    eye_r = torch.eye(self.rank, device=self.U.device, dtype=self.U.dtype)
    loss_u = torch.norm(u_gram - eye_r, p="fro") ** 2
    loss_v = torch.norm(v_gram - eye_r, p="fro") ** 2
    return (loss_u + loss_v) / (self.rank * self.rank)

RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    # Symmetrical square-root initialization
    nn.init.orthogonal_(self.U)
    nn.init.orthogonal_(self.V)
    if self.bias is not None:
        nn.init.zeros_(self.bias)
    if self.norm_weight is not None:
        nn.init.ones_(self.norm_weight)

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None

Expand source code

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
    """Set precomputed U and V low-rank factors."""
    with torch.no_grad():
        self.U.copy_(U_factor)
        self.V.copy_(V_factor)

Set precomputed U and V low-rank factors.

class GalerkinAttention (config: Any | None = None,
rank: int | None = None,
dropout: float = 0.1,
basis_type: str = 'ortho_dct',
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("galerkin", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("galerkin_attention", supports_causal=True, supports_gqa=True)
class GalerkinAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        dropout: float = 0.1,
        basis_type: str = "ortho_dct",
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # ---- Config discovery ----
        if config is None and source_attn is not None:

            class DummyConfig:
                pass

            config = DummyConfig()
            out_layer = getattr(
                source_attn,
                "out_lin",
                getattr(
                    source_attn,
                    "dense",
                    getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                ),
            )
            q_layer = getattr(
                source_attn,
                "q_lin",
                getattr(
                    source_attn,
                    "query",
                    getattr(source_attn, "c_attn", getattr(source_attn, "q_proj", None)),
                ),
            )
            if out_layer is not None and hasattr(out_layer, "out_features"):
                config.hidden_size = out_layer.out_features
            elif q_layer is not None and hasattr(q_layer, "in_features"):
                config.hidden_size = q_layer.in_features
            else:
                config.hidden_size = 768
            config.num_attention_heads = getattr(
                source_attn,
                "n_heads",
                getattr(
                    source_attn,
                    "num_attention_heads",
                    getattr(source_attn, "num_heads", 12),
                ),
            )

        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        self.num_heads: int = (
            getattr(config, "num_attention_heads", getattr(config, "num_heads", 12))
            if config is not None
            else 12
        )
        self.num_kv_heads: int = (
            getattr(config, "num_key_value_heads", self.num_heads)
            if config is not None
            else self.num_heads
        )
        hidden_size: int = (
            getattr(config, "hidden_size", getattr(config, "d_model", 768))
            if config is not None
            else 768
        )
        cfg_head_dim = getattr(config, "head_dim", None) if config is not None else None
        if cfg_head_dim is None and source_attn is not None:
            for attr in ("q_proj", "query", "q_lin", "W_q"):
                mod = getattr(source_attn, attr, None)
                if mod is not None and hasattr(mod, "weight"):
                    cfg_head_dim = mod.weight.shape[0] // max(1, self.num_heads)
                    break
        self.head_dim: int = (
            int(cfg_head_dim)
            if cfg_head_dim is not None
            else getattr(config, "d_kv", hidden_size // max(1, self.num_heads))
        )
        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 = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.rank: int = (
            rank
            if rank is not None
            else getattr(config, "galerkin_rank", getattr(config, "rank", 64))
        )
        self.basis_type: str = basis_type

        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx: int = int(resolved_layer_id)

        # ---- LoRA config ----
        self.lora_rank: int = getattr(
            config,
            "lora_rank",
            getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha: float = getattr(
            config,
            "lora_alpha",
            getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base: bool = self.lora_rank > 0

        # ---- Decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(getattr(config, "model_type", "") or "").lower() if config is not None else ""
            _decoder_prefixes = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            is_dec = False
            if config is not None:
                if getattr(config, "is_decoder", False) or getattr(config, "is_causal", False):
                    is_dec = True
                if any(_mt.startswith(p) for p in _decoder_prefixes):
                    is_dec = True
                if hasattr(config, "architectures") and any(
                    "GPT" in str(a) or "CausalLM" in str(a)
                    for a in (getattr(config, "architectures", []) or [])
                ):
                    is_dec = True
            if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
                is_dec = True
            self.is_decoder = bool(is_dec)
        self.is_causal: bool = bool(self.is_decoder)

        # ---- Nested encoder detection ----
        is_nested = kwargs.get("is_nested_encoder")
        if is_nested is None and source_attn is not None:
            src_cls = source_attn.__class__.__name__
            if (
                src_cls
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                    "ElectraSelfAttention",
                    "DebertaSelfAttention",
                )
                or (src_cls.endswith("SelfAttention") and "distil" not in src_cls.lower())
                or (hasattr(source_attn, "self") and hasattr(source_attn, "output"))
            ):
                is_nested = True
            elif (
                hasattr(source_attn, "out_lin")
                or hasattr(source_attn, "c_proj")
                or hasattr(source_attn, "o_proj")
            ):
                is_nested = False
        if is_nested is None:
            is_nested = False
        self.is_nested_encoder: bool = bool(is_nested)

        # ---- Projections ----
        self.query = LoRALinear(
            hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.key = LoRALinear(
            hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.value = LoRALinear(
            hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", dropout),
                ),
            )
            if config
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)

        self.rotary_emb = None
        self.comp_config: Any | None = comp_config

        # GAL-FIX-07: small cache for the basis matrix
        self._basis_cache: dict[tuple[int, int], torch.Tensor] = {}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    # ---- Weight transfer (unchanged from original, correct) ----
    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")
        if has_out_lin:
            is_nested = False
        elif (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            and hasattr(src, "output")
            and hasattr(src.output, "dense")
        ):
            is_nested = True
            src = src.self
        elif cls_name in (
            "BertSelfAttention",
            "RobertaSelfAttention",
            "XLMRobertaSelfAttention",
            "ElectraSelfAttention",
            "DebertaSelfAttention",
        ) or (cls_name.endswith("SelfAttention") and "distil" not in cls_name.lower()):
            is_nested = True
        else:
            is_nested = False
        self.is_nested_encoder = bool(is_nested)

        if hasattr(src, "rotary_emb"):
            self.rotary_emb = src.rotary_emb

        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
        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()
                cls_c = c_attn_mod.__class__.__name__
                base_cls_c = base_c.__class__.__name__
                if (
                    "Conv1D" in cls_c
                    or "Conv1D" in base_cls_c
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                    or (
                        c_w.ndim == 2
                        and c_w.shape[0] < c_w.shape[1]
                        and "Linear" not in cls_c
                        and "Linear" not in base_cls_c
                    )
                ):
                    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()
                    cls_o = c_proj_mod.__class__.__name__
                    base_cls_o = base_o.__class__.__name__
                    if (
                        "Conv1D" in cls_o
                        or "Conv1D" in base_cls_o
                        or hasattr(c_proj_mod, "nf")
                        or hasattr(base_o, "nf")
                        or (
                            o_w.ndim == 2
                            and o_w.shape[0] < o_w.shape[1]
                            and "Linear" not in cls_o
                            and "Linear" not in base_cls_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
        else:
            proj_candidates = {
                "q": ["q_lin", "q_proj", "query", "W_q"],
                "k": ["k_lin", "k_proj", "key", "W_k"],
                "v": ["v_lin", "v_proj", "value", "W_v"],
                "o": ["out_lin", "out_proj", "dense", "o_proj", "proj", "output"],
            }
            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

        # UniversalAdapter fallback
        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 warning: %s", e)

        target_map = {
            "q": self.query,
            "k": self.key,
            "v": self.value,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                w_norm = float(torch.linalg.norm(target_mod.weight.data.float()).item())
                if w_norm < 1e-4:
                    logger.warning("GalerkinAttention: %s near-zero norm (%.2e)", proj_key, w_norm)

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    # ------------------------------------------------------------------
    # Basis construction (GAL-FIX-07: cached)
    # ------------------------------------------------------------------
    def _get_orthonormal_basis(
        self,
        N: int,
        r: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        r_actual = max(1, min(r, N))
        key = (N, r_actual)
        cached = self._basis_cache.get(key)
        if cached is not None and cached.device == device and cached.dtype == dtype:
            return cached

        n = torch.arange(N, device=device, dtype=torch.float32).unsqueeze(1)
        k = torch.arange(r_actual, device=device, dtype=torch.float32).unsqueeze(0)

        phi = torch.cos((math.pi * k * (2.0 * n + 1.0)) / (2.0 * max(1, N)))
        phi[:, 0] = phi[:, 0] * (1.0 / math.sqrt(max(1, N)))
        if r_actual > 1:
            phi[:, 1:] = phi[:, 1:] * math.sqrt(2.0 / max(1, N))

        eye = torch.eye(N, r_actual, device=device, dtype=torch.float32)
        phi = phi + 1e-4 * eye
        phi, _ = torch.linalg.qr(phi)
        phi = phi.to(dtype=dtype)
        self._basis_cache[key] = phi
        return phi

    # ------------------------------------------------------------------
    # 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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        is_cross_attention = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross_attention else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Linear projections
        Q = self.query(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous()
        K = self.key(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V = self.value(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # 2. RoPE (GAL-FIX-05)
        if position_embeddings is not None and not is_cross_attention:
            if isinstance(position_embeddings, (tuple, list)):
                cos, sin = position_embeddings[0], position_embeddings[1]
            else:
                cos, sin = position_embeddings, None
            if sin is not None:
                Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids)

        # 3. GQA
        if self.num_kv_heads != self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        eff_mask = encoder_attention_mask if is_cross_attention else attention_mask

        # 4. Decoder causal path
        if self.is_decoder and not is_cross_attention:
            # GAL-FIX-03: only use SDPA when we don't need attn weights/head_mask
            if not output_attentions and head_mask is None:
                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 = N > 1
                    _sdpa_out = F.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)
                    _sdpa_out = (
                        _sdpa_out.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
                    )
                    _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                    out = self.out_proj(_sdpa_out) if not self.is_nested_encoder else _sdpa_out
                    return out, None
                except Exception:
                    pass

            # Manual causal path
            scale = d**-0.5
            scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

            # GAL-FIX-04: explicit causal mask
            causal_diag = 1 + max(0, kv_len - N)
            causal_mask = torch.triu(
                torch.full((N, kv_len), float("-inf"), device=device, dtype=torch.float32),
                diagonal=causal_diag,
            )
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

            if eff_mask is not None:
                if eff_mask.dim() == 2:
                    pad_m = (1.0 - eff_mask.unsqueeze(1).unsqueeze(2).float()) * -1e4
                    scores = scores + pad_m
                elif eff_mask.dim() == 3:
                    scores = scores + eff_mask.unsqueeze(1).float()
                elif eff_mask.dim() == 4:
                    scores = scores + eff_mask.float()

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            scores_max = scores.max(dim=-1, keepdim=True).values.detach()
            exp_scores = torch.exp((scores - scores_max).clamp(min=-50.0, max=0.0))
            denom = exp_scores.sum(dim=-1, keepdim=True).clamp(min=1e-6)
            probs = (exp_scores / denom).to(dtype)

            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)
                probs = probs * head_mask.to(probs.dtype)

            probs = self.dropout(probs)
            context = torch.matmul(probs, V)
            attn_weights = probs if output_attentions else None

        # 5. Encoder: softmax in basis space (GAL-FIX-01: original formulation)
        else:
            valid_mask = get_token_valid_mask(eff_mask, B, kv_len, device)
            if valid_mask is None:
                valid_mask = torch.ones(B, kv_len, device=device, dtype=torch.bool)
            mask_expanded = valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)

            K_masked = K * mask_expanded
            V_masked = V * mask_expanded

            r_actual = max(1, min(self.rank, kv_len))
            Phi = self._get_orthonormal_basis(kv_len, r_actual, device, dtype)
            # Phi: (N, r)

            # Project into basis
            Q_basis = torch.einsum("bhnd,nr->bhrd", Q, Phi)  # (B, H, r, d)
            K_basis = torch.einsum("bhnd,nr->bhrd", K_masked, Phi)
            V_basis = torch.einsum("bhnd,nr->bhrd", V_masked, Phi)

            scale = d**-0.5
            scores = (
                torch.matmul(Q_basis.float(), K_basis.float().transpose(-2, -1).contiguous())
                * scale
            )
            # (B, H, r, r)

            # Numerically stable softmax in basis space
            scores_max = scores.max(dim=-1, keepdim=True).values.detach()
            exp_scores = torch.exp((scores - scores_max).clamp(min=-50.0, max=0.0))
            denom = exp_scores.sum(dim=-1, keepdim=True).clamp(min=1e-6)
            probs = exp_scores / denom
            probs = torch.nan_to_num(probs, nan=0.0)

            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)
                probs = probs * head_mask.to(probs.dtype)

            probs = self.dropout(probs)

            # Attention in basis space
            O_basis = torch.matmul(probs.to(dtype), V_basis)  # (B, H, r, d)

            # Reconstruct sequence representation
            context = torch.einsum("bhrd,nr->bhnd", O_basis, Phi)  # (B, H, N, d)

            # Apply padding mask on outputs
            if eff_mask is not None and not is_cross_attention:
                context = context * mask_expanded

            attn_weights = probs if output_attentions else None

        # 6. Reshape and output projection
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

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

        return out, attn_weights

Every 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.

Instance variables

prop W_k

Expand source code

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

prop W_q

Expand source code

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

prop W_v

Expand source code

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

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.key

prop k_proj

Expand source code

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

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.query

prop q_proj

Expand source code

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

prop v_lin

Expand source code

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

prop v_proj

Expand source code

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

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class GatedDeltaNet (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
num_heads: int | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
old_attn: nn.Module | None = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "gated_delta_net",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class GatedDeltaNet(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        comp_config: Any | None = None,
        layer_id: int | None = None,
        layer_idx: int | None = None,
        old_attn: nn.Module | None = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # ---- d_model ----
        if d_model is None:
            d_model = kwargs.get("d_model") or kwargs.get("hidden_size") or kwargs.get("dim")
            if d_model is None:
                if config is not None:
                    d_model = getattr(config, "hidden_size", getattr(config, "dim", 768))
                elif source_attn is not None:
                    out_layer = None
                    for attr in ("out_lin", "dense", "c_proj", "o_proj", "out_proj"):
                        if hasattr(source_attn, attr):
                            out_layer = getattr(source_attn, attr)
                            break
                    if out_layer is not None:
                        d_model = getattr(
                            out_layer,
                            "out_features",
                            getattr(out_layer, "in_features", 768),
                        )
                    else:
                        for attr in ("q_proj", "query", "q_lin", "W_q"):
                            q_l = getattr(source_attn, attr, None)
                            if q_l is not None:
                                d_model = getattr(
                                    q_l,
                                    "in_features",
                                    getattr(q_l, "out_features", 768),
                                )
                                break
                if d_model is None:
                    d_model = 768
        if hasattr(d_model, "hidden_size"):
            d_model = d_model.hidden_size

        # ---- num_heads ----
        if num_heads is None:
            num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")
            if num_heads is None:
                if config is not None:
                    num_heads = getattr(
                        config,
                        "num_attention_heads",
                        getattr(config, "n_heads", getattr(config, "num_heads", 12)),
                    )
                elif source_attn is not None:
                    num_heads = getattr(
                        source_attn,
                        "n_heads",
                        getattr(
                            source_attn,
                            "num_attention_heads",
                            getattr(source_attn, "num_heads", 12),
                        ),
                    )
                else:
                    num_heads = 12

        # GDN-FIX-01: None-safe
        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        self.d_model = int(d_model)
        self.num_heads = max(1, int(num_heads))

        # ---- num_kv_heads ----
        cfg_kv = None
        if config is not None:
            cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        if cfg_kv is None:
            cfg_kv = kwargs.get("num_kv_heads", self.num_heads)
        self.num_kv_heads = max(1, min(int(cfg_kv), self.num_heads))

        # ---- head_dim ----
        head_dim = getattr(config, "head_dim", None) if config else None
        if head_dim is None and source_attn is not None:
            for q_cand in ("q_proj", "query", "q_lin", "W_q"):
                q_mod = getattr(source_attn, q_cand, None)
                if q_mod is not None and hasattr(q_mod, "out_features"):
                    head_dim = q_mod.out_features // max(1, self.num_heads)
                    break
        if head_dim is None:
            head_dim = self.d_model // self.num_heads
        self.head_dim = int(head_dim)
        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 // self.num_kv_heads)

        resolved_layer_id = (
            layer_id
            if layer_id is not None
            else (
                layer_idx
                if layer_idx is not None
                else kwargs.get("layer_id", kwargs.get("layer_idx", 0))
            )
        )
        self.layer_idx = int(resolved_layer_id)

        # ---- LoRA ----
        if config is not None:
            self.lora_rank = int(
                getattr(
                    config,
                    "lora_rank",
                    getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
                )
            )
            self.lora_alpha = float(
                getattr(
                    config,
                    "lora_alpha",
                    getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
                )
            )
        elif comp_config is not None:
            self.lora_rank = int(getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)))
            self.lora_alpha = float(
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
            )
        else:
            self.lora_rank = int(kwargs.get("lora_rank", 16))
            self.lora_alpha = float(kwargs.get("lora_alpha", 32.0))
        self.freeze_base = self.lora_rank > 0

        # ---- decoder detection ----
        model_type = str(getattr(config, "model_type", "") if config else "").lower()
        is_dec = False
        if config is not None:
            is_dec = bool(
                getattr(config, "is_decoder", False) or getattr(config, "is_causal", False)
            )
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
            "codegen",
        )
        if any(model_type.startswith(p) for p in decoder_prefixes):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        if source_attn is not None:
            if getattr(source_attn, "is_decoder", False) or hasattr(source_attn, "c_attn"):
                is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = self.is_decoder

        # ---- nested encoder detection ----
        is_nested = False
        if source_attn is not None:
            src_cls = source_attn.__class__.__name__
            if (
                (
                    src_cls
                    in (
                        "BertSelfAttention",
                        "RobertaSelfAttention",
                        "XLMRobertaSelfAttention",
                        "ElectraSelfAttention",
                        "DebertaSelfAttention",
                    )
                    or (src_cls.endswith("SelfAttention") and "distil" not in src_cls.lower())
                )
                or hasattr(source_attn, "self")
                and hasattr(source_attn, "output")
            ):
                is_nested = True
        if (
            not is_nested
            and model_type in ("roberta", "bert", "xlm-roberta")
            and not self.is_decoder
        ):
            is_nested = True
        self.is_nested_encoder = bool(is_nested)

        # ---- projections ----
        self.q_proj = LoRALinear(
            self.d_model,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.d_model,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.d_model,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.z_proj = LoRALinear(
            self.d_model,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.d_model,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # GDN-FIX-13: single output scale
        self.output_scale = nn.Parameter(torch.ones(1))

        # ---- gating ----
        self.beta_proj = nn.Linear(self.d_model, self.num_heads, bias=True)
        self.alpha_proj = nn.Linear(self.d_model, self.num_heads, bias=True)
        with torch.no_grad():
            nn.init.constant_(self.beta_proj.bias, 2.0)  # ~0.88 initial decay
            nn.init.constant_(self.alpha_proj.bias, 0.0)  # ~0.5 initial step

        # ---- chunk size (GDN-FIX-14) ----
        cs = int(_cfg_get(comp_config, "chunk_size", kwargs.get("chunk_size", 32)))
        self.chunk_size = max(8, min(cs, 128))

        self.source_attn = source_attn
        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )
        self.comp_config = comp_config

        # ---- storage accounting ----
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

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

    def freeze_base_weights(self):
        for m in (self.q_proj, self.k_proj, self.v_proj, self.z_proj, self.out_proj):
            if hasattr(m, "freeze_base_weights"):
                m.freeze_base_weights()

    # ---- aliases ----
    @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 o_proj(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 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 c_proj(self):
        return self.out_proj

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

        try:
            for p in old_attn.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        cls_name = old_attn.__class__.__name__
        is_nested = False
        if (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
                "DebertaSelfAttention",
            )
            or (cls_name.endswith("SelfAttention") and "distil" not in cls_name.lower())
        ):
            is_nested = True
        if not is_nested and getattr(self, "is_nested_encoder", False):
            is_nested = True
        self.is_nested_encoder = bool(is_nested)

        if hasattr(old_attn, "rotary_emb"):
            self.rotary_emb = old_attn.rotary_emb

        extracted = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused c_attn
        if hasattr(old_attn, "c_attn") and old_attn.c_attn is not None:
            c = old_attn.c_attn
            # GDN-FIX-12: guarded merge
            if hasattr(c, "merge") and callable(c.merge) and not getattr(c, "merged", False):
                try:
                    c.merge()
                except Exception as e:
                    logger.debug("c_attn merge skipped: %s", e)
            base_c = getattr(c, "base_layer", c)
            cw = getattr(base_c, "weight", getattr(c, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                cls_c = c.__class__.__name__
                base_cls_c = base_c.__class__.__name__
                is_conv = (
                    "Conv1D" in cls_c
                    or "Conv1D" in base_cls_c
                    or hasattr(c, "nf")
                    or hasattr(base_c, "nf")
                    or (cw.ndim == 2 and cw.shape[0] < cw.shape[1])
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )

            if hasattr(old_attn, "c_proj") and old_attn.c_proj is not None:
                o = old_attn.c_proj
                base_o = getattr(o, "base_layer", o)
                ow = getattr(base_o, "weight", getattr(o, "weight", None))
                if ow is not None:
                    ow = ow.data.clone()
                    is_conv_o = (
                        "Conv1D" in o.__class__.__name__
                        or "Conv1D" in base_o.__class__.__name__
                        or hasattr(o, "nf")
                        or hasattr(base_o, "nf")
                        or (ow.ndim == 2 and ow.shape[0] < ow.shape[1])
                    )
                    if is_conv_o:
                        ow = ow.t().contiguous()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # Nested (BERT / RoBERTa)
        elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
            s = old_attn.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(old_attn, "output") and hasattr(old_attn.output, "dense"):
                d = old_attn.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat
        else:
            candidates = {
                "q": ("q_lin", "q_proj", "query", "W_q"),
                "k": ("k_lin", "k_proj", "key", "W_k"),
                "v": ("v_lin", "v_proj", "value", "W_v"),
                "o": ("out_lin", "out_proj", "dense", "o_proj", "output", "proj"),
            }
            for key, names in candidates.items():
                for n in names:
                    if hasattr(old_attn, n):
                        m = getattr(old_attn, n)
                        if m is not None and hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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(old_attn, 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:
            logger.warning("GatedDeltaNet copy_weights: missing %s", missing_after)
            return

        with torch.no_grad():
            for key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target.weight)
                        if target.bias is not None:
                            nn.init.zeros_(target.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target.weight.shape
                    and w.t().contiguous().shape == target.weight.shape
                ):
                    w = w.t().contiguous()

                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))
                target.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                )

                if b is not None and target.bias is not None:
                    mb = min(target.bias.numel(), b.numel())
                    target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))
                elif target.bias is not None:
                    nn.init.zeros_(target.bias)

                n = float(torch.linalg.norm(target.weight.data.float()).item())
                if n < 1e-4:
                    logger.warning("GatedDeltaNet: '%s' near-zero norm (%.2e)", key, n)

                if self.freeze_base:
                    target.freeze_base_weights()

            # z_proj from q
            w_q = extracted["q"][0]
            if w_q is not None:
                out_z = min(self.z_proj.weight.size(0), w_q.size(0))
                in_z = min(self.z_proj.weight.size(1), w_q.size(1))
                self.z_proj.weight.data[:out_z, :in_z].copy_(
                    w_q[:out_z, :in_z].to(self.z_proj.weight.device, self.z_proj.weight.dtype)
                )
                if self.freeze_base:
                    self.z_proj.freeze_base_weights()

        try:
            for m in (
                self.q_proj,
                self.k_proj,
                self.v_proj,
                self.z_proj,
                self.out_proj,
            ):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        if self.freeze_base:
            self.freeze_base_weights()

    # ------------------------------------------------------------------
    # Mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_valid_mask(mask, B, L, device):
        """GDN-FIX-04: returns (B, L) bool or None."""
        if mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if mask.dim() == 2:
            m = _to_bool(mask)
            if m.shape == (B, L):
                return m
        elif mask.dim() == 3:
            m = _to_bool(mask)
            if m.shape[0] == B and m.shape[-1] == L:
                return m[:, -1, :]
        elif mask.dim() == 4:
            m = _to_bool(mask)
            m = m.any(dim=-2)
            if m.dim() == 3:
                m = m.squeeze(1)
            if m.shape == (B, L):
                return m
        return None

    # ------------------------------------------------------------------
    # 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_ids: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ):
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim

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

        q = self.q_proj(hidden_states).view(B, L, H, d).transpose(1, 2).contiguous()
        k = self.k_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()
        v = self.v_proj(kv_input).view(B, kv_len, self.num_kv_heads, d).transpose(1, 2).contiguous()
        z = self.z_proj(hidden_states)

        # ---- RoPE ----
        if position_embeddings is not None and not is_cross:
            if isinstance(position_embeddings, tuple):
                cos, sin = position_embeddings
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
        elif self.rotary_emb is not None and not is_cross:
            if position_ids is None:
                position_ids = (
                    torch.arange(kv_len, device=hidden_states.device).unsqueeze(0).expand(B, -1)
                )
            try:
                cos, sin = self.rotary_emb(k, seq_len=kv_len)
            except TypeError:
                try:
                    cos, sin = self.rotary_emb(k, position_ids)
                except Exception:
                    cos, sin = None, None
            if cos is not None and sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)

        # ---- GQA ----
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        # ---- Gating ----
        hs_beta = hidden_states.to(self.beta_proj.weight.dtype)
        beta = torch.sigmoid(self.beta_proj(hs_beta)).transpose(1, 2).contiguous()  # (B,H,L)
        hs_alpha = hidden_states.to(self.alpha_proj.weight.dtype)
        alpha = torch.sigmoid(self.alpha_proj(hs_alpha)).transpose(1, 2).contiguous()

        # ---- Mask ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        mask_bool = self._normalize_valid_mask(eff_mask, B, kv_len, hidden_states.device)

        # ---- is_causal (GDN-FIX-08) ----
        bool(self.is_decoder and not is_cross and L > 1 and past_key_value is None)

        # ---- Delta scan ----
        y_fwd = chunked_delta_rule_scan(
            q,
            k,
            v,
            beta,
            alpha,
            chunk_size=self.chunk_size,
        )
        if not self.is_decoder:
            q_b = torch.flip(q, dims=[2])
            k_b = torch.flip(k, dims=[2])
            v_b = torch.flip(v, dims=[2])
            beta_b = torch.flip(beta, dims=[2])
            alpha_b = torch.flip(alpha, dims=[2])
            y_bwd_f = chunked_delta_rule_scan(
                q_b, k_b, v_b, beta_b, alpha_b, chunk_size=self.chunk_size
            )
            y_bwd = torch.flip(y_bwd_f, dims=[2])
            y = 0.5 * (y_fwd + y_bwd)
        else:
            y = y_fwd

        # ---- head_mask (GDN-FIX-11) ----
        if head_mask is not None:
            hm = head_mask.to(y.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            y = y * hm

        # ---- Recombine + SiLU gating ----
        y = torch.nan_to_num(y, nan=0.0, posinf=1e4, neginf=-1e4)
        y = y.transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        gated_y = y.to(z.dtype) * F.silu(z)

        # ---- Output projection ----
        if self.is_nested_encoder:
            out = gated_y
        else:
            out = self.out_proj(gated_y)

        out = out * self.output_scale

        # ---- mask output (only once; GDN-FIX-05) ----
        if mask_bool is not None:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        return out, None

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("q_proj", self.q_proj),
            ("k_proj", self.k_proj),
            ("v_proj", self.v_proj),
            ("z_proj", self.z_proj),
            ("out_proj", self.out_proj),
            ("beta_proj", self.beta_proj),
            ("alpha_proj", self.alpha_proj),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["output_scale"] = _sz(self.output_scale)
        out["total"] = sum(out.values())
        return out

Every 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.

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_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 freeze_base_weights(self)

Expand source code

def freeze_base_weights(self):
    for m in (self.q_proj, self.k_proj, self.v_proj, self.z_proj, self.out_proj):
        if hasattr(m, "freeze_base_weights"):
            m.freeze_base_weights()

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("q_proj", self.q_proj),
        ("k_proj", self.k_proj),
        ("v_proj", self.v_proj),
        ("z_proj", self.z_proj),
        ("out_proj", self.out_proj),
        ("beta_proj", self.beta_proj),
        ("alpha_proj", self.alpha_proj),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["output_scale"] = _sz(self.output_scale)
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class GatedDeltaNetBlock (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
num_heads: int | None = None,
use_full_attention: bool = False,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("gated_deltanet_block")
class GatedDeltaNetBlock(BaseAttention):
    r"""Gated DeltaNet Transformer Block with Pretrained Weight Transfer.

    Mathematical Formulation:
        $$\text{Output} = \begin{cases}
            \text{LayerNorm}(X + \text{MHA}(X, X, X)), & \text{if } \text{use\_full\_attention} \\
            \text{LayerNorm}(X + \text{GatedDeltaNet}(X)), & \text{otherwise}
        \end{cases}$$
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        use_full_attention: bool = False,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        if d_model is None:
            if config is not None:
                d_model = getattr(config, "hidden_size", getattr(config, "dim", 768))
            else:
                d_model = kwargs.get("hidden_size", 768)

        if num_heads is None:
            if config is not None:
                num_heads = getattr(config, "num_attention_heads", 4)
            else:
                num_heads = kwargs.get("num_attention_heads", 4)

        self.d_model: int = d_model
        self.num_heads: int = num_heads
        self.use_full_attention: bool = use_full_attention
        self.comp_config = comp_config

        if use_full_attention:
            self.attn: nn.MultiheadAttention | None = nn.MultiheadAttention(
                self.d_model, self.num_heads, batch_first=True
            )
        else:
            self.attn = None

        self.deltanet = GatedDeltaNet(
            config=config,
            comp_config=comp_config,
            d_model=self.d_model,
            num_heads=self.num_heads,
            **kwargs,
        )
        self.norm = nn.LayerNorm(self.d_model)
        self.is_nested_encoder: bool = getattr(self.deltanet, "is_nested_encoder", False)

    # Compatibility properties
    @property
    def query(self) -> Any:
        return self.deltanet.query

    @property
    def key(self) -> Any:
        return self.deltanet.key

    @property
    def value(self) -> Any:
        return self.deltanet.value

    @property
    def q_proj(self) -> Any:
        return self.deltanet.q_proj

    @property
    def k_proj(self) -> Any:
        return self.deltanet.k_proj

    @property
    def v_proj(self) -> Any:
        return self.deltanet.v_proj

    @property
    def out_proj(self) -> Any:
        return self.deltanet.out_proj

    @property
    def q_lin(self) -> Any:
        return self.deltanet.q_lin

    @property
    def k_lin(self) -> Any:
        return self.deltanet.k_lin

    @property
    def v_lin(self) -> Any:
        return self.deltanet.v_lin

    @property
    def out_lin(self) -> Any:
        return self.deltanet.out_lin

    def copy_weights(self, old_attn: Any) -> None:
        """Transfer pretrained weights to either MHA or GatedDeltaNet."""
        if self.use_full_attention and self.attn is not None:
            try:
                adapter = UniversalAdapter()
                w_q, b_q = adapter.get_weight_bias(old_attn, "q")
                w_k, b_k = adapter.get_weight_bias(old_attn, "k")
                w_v, b_v = adapter.get_weight_bias(old_attn, "v")
                w_o, b_o = adapter.get_weight_bias(old_attn, "o")

                if w_o is None:
                    for cand in ("out_proj", "out_lin", "o_proj", "output", "proj"):
                        for target in (old_attn, getattr(old_attn, "self", None)):
                            if target is not None and hasattr(target, cand):
                                mod = getattr(target, cand)
                                if hasattr(mod, "dense"):
                                    mod = mod.dense
                                if hasattr(mod, "weight") and mod.weight is not None:
                                    w_o = mod.weight.data.clone()
                                    if w_o.ndim == 2 and w_o.shape[0] < w_o.shape[1]:
                                        w_o = w_o.t()
                                    b_o = getattr(mod, "bias", None)
                                    b_o = b_o.data.clone() if b_o is not None else None
                                    break
                        if w_o is not None:
                            break

                with torch.no_grad():
                    if w_q is not None and w_k is not None and w_v is not None:
                        qkv_w = torch.cat([w_q, w_k, w_v], dim=0)
                        if qkv_w.shape == self.attn.in_proj_weight.shape:
                            self.attn.in_proj_weight.copy_(qkv_w)
                        if (
                            b_q is not None
                            and b_k is not None
                            and b_v is not None
                            and self.attn.in_proj_bias is not None
                        ):
                            cat_b = torch.cat([b_q, b_k, b_v], dim=0)
                            if cat_b.shape == self.attn.in_proj_bias.shape:
                                self.attn.in_proj_bias.copy_(cat_b)
                    if w_o is not None:
                        if w_o.shape == self.attn.out_proj.weight.shape:
                            self.attn.out_proj.weight.copy_(w_o)
                        elif w_o.t().shape == self.attn.out_proj.weight.shape:
                            self.attn.out_proj.weight.copy_(w_o.t())
            except (AttributeError, KeyError, RuntimeError, TypeError, ValueError) as e:
                logger.warning(
                    "GatedDeltaNetBlock copy_weights to MultiheadAttention incomplete: %s",
                    e,
                )
        else:
            if hasattr(self.deltanet, "copy_weights"):
                self.deltanet.copy_weights(old_attn)
                self.is_nested_encoder = getattr(self.deltanet, "is_nested_encoder", False)

    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]:
        """Forward pass for GatedDeltaNetBlock."""
        x = hidden_states
        if self.use_full_attention and self.attn is not None:
            key_padding_mask = None
            if attention_mask is not None:
                if attention_mask.dim() == 2:
                    key_padding_mask = attention_mask == 0
                elif attention_mask.dim() == 4:
                    key_padding_mask = attention_mask[:, 0, 0, :] <= -10000.0

            attn_out, attn_weights = self.attn(
                x,
                x,
                x,
                key_padding_mask=key_padding_mask,
                need_weights=output_attentions,
            )
            if attn_out.shape[-1] != x.shape[-1]:
                attn_out = (
                    F.pad(attn_out, (0, x.shape[-1] - attn_out.shape[-1]))
                    if attn_out.shape[-1] < x.shape[-1]
                    else attn_out[..., : x.shape[-1]]
                )
            if self.is_nested_encoder:
                return attn_out, (attn_weights if output_attentions else None)
            out = self.norm(x + attn_out)
            return out, (attn_weights if output_attentions else None)
        else:
            delta_out, _ = self.deltanet(
                x,
                attention_mask=attention_mask,
                head_mask=head_mask,
                output_attentions=output_attentions,
                **kwargs,
            )
            if delta_out.shape[-1] != x.shape[-1]:
                delta_out = (
                    F.pad(delta_out, (0, x.shape[-1] - delta_out.shape[-1]))
                    if delta_out.shape[-1] < x.shape[-1]
                    else delta_out[..., : x.shape[-1]]
                )
            if self.is_nested_encoder:
                return delta_out, None
            if getattr(self, "is_nested_encoder", False):
                return delta_out, None
            out = self.norm(x + delta_out)
            return out, None

Gated DeltaNet Transformer Block with Pretrained Weight Transfer.

Mathematical Formulation: $$\text{Output} = \begin{cases} \text{LayerNorm}(X + \text{MHA}(X, X, X)), & \text{if } \text{use_full_attention} \ \text{LayerNorm}(X + \text{GatedDeltaNet}(X)), & \text{otherwise} \end{cases}$$

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.

Instance variables

prop k_lin : Any

Expand source code

@property
def k_lin(self) -> Any:
    return self.deltanet.k_lin

prop k_proj : Any

Expand source code

@property
def k_proj(self) -> Any:
    return self.deltanet.k_proj

prop key : Any

Expand source code

@property
def key(self) -> Any:
    return self.deltanet.key

prop out_lin : Any

Expand source code

@property
def out_lin(self) -> Any:
    return self.deltanet.out_lin

prop out_proj : Any

Expand source code

@property
def out_proj(self) -> Any:
    return self.deltanet.out_proj

prop q_lin : Any

Expand source code

@property
def q_lin(self) -> Any:
    return self.deltanet.q_lin

prop q_proj : Any

Expand source code

@property
def q_proj(self) -> Any:
    return self.deltanet.q_proj

prop query : Any

Expand source code

@property
def query(self) -> Any:
    return self.deltanet.query

prop v_lin : Any

Expand source code

@property
def v_lin(self) -> Any:
    return self.deltanet.v_lin

prop v_proj : Any

Expand source code

@property
def v_proj(self) -> Any:
    return self.deltanet.v_proj

prop value : Any

Expand source code

@property
def value(self) -> Any:
    return self.deltanet.value

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Transfer pretrained weights to either MHA or GatedDeltaNet."""
    if self.use_full_attention and self.attn is not None:
        try:
            adapter = UniversalAdapter()
            w_q, b_q = adapter.get_weight_bias(old_attn, "q")
            w_k, b_k = adapter.get_weight_bias(old_attn, "k")
            w_v, b_v = adapter.get_weight_bias(old_attn, "v")
            w_o, b_o = adapter.get_weight_bias(old_attn, "o")

            if w_o is None:
                for cand in ("out_proj", "out_lin", "o_proj", "output", "proj"):
                    for target in (old_attn, getattr(old_attn, "self", None)):
                        if target is not None and hasattr(target, cand):
                            mod = getattr(target, cand)
                            if hasattr(mod, "dense"):
                                mod = mod.dense
                            if hasattr(mod, "weight") and mod.weight is not None:
                                w_o = mod.weight.data.clone()
                                if w_o.ndim == 2 and w_o.shape[0] < w_o.shape[1]:
                                    w_o = w_o.t()
                                b_o = getattr(mod, "bias", None)
                                b_o = b_o.data.clone() if b_o is not None else None
                                break
                    if w_o is not None:
                        break

            with torch.no_grad():
                if w_q is not None and w_k is not None and w_v is not None:
                    qkv_w = torch.cat([w_q, w_k, w_v], dim=0)
                    if qkv_w.shape == self.attn.in_proj_weight.shape:
                        self.attn.in_proj_weight.copy_(qkv_w)
                    if (
                        b_q is not None
                        and b_k is not None
                        and b_v is not None
                        and self.attn.in_proj_bias is not None
                    ):
                        cat_b = torch.cat([b_q, b_k, b_v], dim=0)
                        if cat_b.shape == self.attn.in_proj_bias.shape:
                            self.attn.in_proj_bias.copy_(cat_b)
                if w_o is not None:
                    if w_o.shape == self.attn.out_proj.weight.shape:
                        self.attn.out_proj.weight.copy_(w_o)
                    elif w_o.t().shape == self.attn.out_proj.weight.shape:
                        self.attn.out_proj.weight.copy_(w_o.t())
        except (AttributeError, KeyError, RuntimeError, TypeError, ValueError) as e:
            logger.warning(
                "GatedDeltaNetBlock copy_weights to MultiheadAttention incomplete: %s",
                e,
            )
    else:
        if hasattr(self.deltanet, "copy_weights"):
            self.deltanet.copy_weights(old_attn)
            self.is_nested_encoder = getattr(self.deltanet, "is_nested_encoder", False)

Transfer pretrained weights to either MHA or GatedDeltaNet.

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]

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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Forward pass for GatedDeltaNetBlock."""
    x = hidden_states
    if self.use_full_attention and self.attn is not None:
        key_padding_mask = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                key_padding_mask = attention_mask == 0
            elif attention_mask.dim() == 4:
                key_padding_mask = attention_mask[:, 0, 0, :] <= -10000.0

        attn_out, attn_weights = self.attn(
            x,
            x,
            x,
            key_padding_mask=key_padding_mask,
            need_weights=output_attentions,
        )
        if attn_out.shape[-1] != x.shape[-1]:
            attn_out = (
                F.pad(attn_out, (0, x.shape[-1] - attn_out.shape[-1]))
                if attn_out.shape[-1] < x.shape[-1]
                else attn_out[..., : x.shape[-1]]
            )
        if self.is_nested_encoder:
            return attn_out, (attn_weights if output_attentions else None)
        out = self.norm(x + attn_out)
        return out, (attn_weights if output_attentions else None)
    else:
        delta_out, _ = self.deltanet(
            x,
            attention_mask=attention_mask,
            head_mask=head_mask,
            output_attentions=output_attentions,
            **kwargs,
        )
        if delta_out.shape[-1] != x.shape[-1]:
            delta_out = (
                F.pad(delta_out, (0, x.shape[-1] - delta_out.shape[-1]))
                if delta_out.shape[-1] < x.shape[-1]
                else delta_out[..., : x.shape[-1]]
            )
        if self.is_nested_encoder:
            return delta_out, None
        if getattr(self, "is_nested_encoder", False):
            return delta_out, None
        out = self.norm(x + delta_out)
        return out, None

Forward pass for GatedDeltaNetBlock.

Inherited members

• BaseAttention:
  • extra_repr

class GatedLinearAttention (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
num_heads: int | None = None,
head_dim: int | None = None,
num_kv_heads: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
chunk_size: int = 128,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("gated_linear_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("gated_linear_attn", supports_causal=True, supports_gqa=True)
class GatedLinearAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "linear_attn"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        head_dim: int | None = None,
        num_kv_heads: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        chunk_size: int = 128,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        src_attn = kwargs.get("pretrained_attn") or kwargs.get("old_attn")

        # Merge explicit args into kwargs for resolver
        _kw = dict(kwargs)
        if d_model is not None:
            _kw["d_model"] = (
                d_model if isinstance(d_model, int) else getattr(d_model, "hidden_size", None)
            )
        if num_heads is not None:
            _kw["num_heads"] = num_heads
        if head_dim is not None:
            _kw["head_dim"] = head_dim
        if num_kv_heads is not None:
            _kw["num_kv_heads"] = num_kv_heads

        # GLA3-FIX-02 / GLA3-FIX-04
        resolved_dim, resolved_heads, resolved_head_dim, resolved_kv_heads = _resolve_dim_and_heads(
            config, src_attn, _kw
        )

        self.dim = resolved_dim
        self.hidden_size = self.dim
        self.num_heads = resolved_heads
        self.head_dim = resolved_head_dim
        self.num_kv_heads = resolved_kv_heads

        # GLA3-FIX-01: hard consistency check
        if self.num_heads <= 0:
            raise ValueError(f"num_heads must be > 0, got {self.num_heads}")
        if self.head_dim <= 0:
            raise ValueError(f"head_dim must be > 0, got {self.head_dim}")
        if self.dim != self.num_heads * self.head_dim:
            # This is the smoking gun for the 12-vs-16 error
            raise ValueError(
                f"[GLA] Dimension mismatch in config: "
                f"hidden_size={self.dim}, num_heads={self.num_heads}, "
                f"head_dim={self.head_dim}, "
                f"num_heads * head_dim = {self.num_heads * self.head_dim} != {self.dim}. "
                f"Fix by passing explicit head_dim or set config.head_dim."
            )

        self.all_head_size = self.num_heads * self.head_dim
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.kv_head_size = self.num_kv_heads * self.head_dim

        self.chunk_size = max(16, int(chunk_size))
        self.lora_rank = int(lora_rank)
        self.lora_alpha = float(lora_alpha)

        self.is_decoder = _detect_decoder(config, is_decoder, src_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # ---- Projections ----
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.z_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        # GLA3-FIX-03: build decay_gate/log_decay from the same `num_heads`
        self.decay_gate = nn.Linear(self.dim, self.num_heads, bias=True)
        with torch.no_grad():
            nn.init.zeros_(self.decay_gate.weight)
            nn.init.zeros_(self.decay_gate.bias)

        self.log_decay = nn.Parameter(torch.zeros(self.num_heads))

        self.state_norm = nn.LayerNorm(self.head_dim)
        self.output_scale = nn.Parameter(torch.ones(1))

        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        # Sanity: ensure z_proj output matches all_head_size
        if self.z_proj.out_features != self.all_head_size:
            raise RuntimeError(
                f"GLA: z_proj.out_features={self.z_proj.out_features} "
                f"!= all_head_size={self.all_head_size}"
            )

        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0
        self.triton_kwargs = {"phi": "elu+1", "eps": 1e-4}

        if src_attn is not None:
            self.copy_weights(src_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 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 o_proj(self):
        return self.out_proj

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

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

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

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

        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        cls_name = src.__class__.__name__
        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())
        ):
            is_nested = True
        self.is_nested_encoder = bool(is_nested)

        extracted = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c = src.c_attn
            base_c = getattr(c, "base_layer", c)
            cw = getattr(base_c, "weight", getattr(c, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                is_conv = (
                    c.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or hasattr(c, "nf")
                    or hasattr(base_c, "nf")
                    or (cw.ndim == 2 and cw.shape[0] < cw.shape[1])
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )

            if hasattr(src, "c_proj") and src.c_proj is not None:
                o = src.c_proj
                base_o = getattr(o, "base_layer", o)
                ow = getattr(base_o, "weight", getattr(o, "weight", None))
                if ow is not None:
                    ow = ow.data.clone()
                    is_conv_o = (
                        o.__class__.__name__ == "Conv1D"
                        or base_o.__class__.__name__ == "Conv1D"
                        or hasattr(o, "nf")
                        or hasattr(base_o, "nf")
                        or (ow.ndim == 2 and ow.shape[0] < ow.shape[1])
                    )
                    if is_conv_o:
                        ow = ow.t().contiguous()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # Nested BERT / RoBERTa
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                d = src.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat (DistilBERT, Qwen, LLaMA)
        else:
            candidates = {
                "q": ("q_proj", "q_lin", "query"),
                "k": ("k_proj", "k_lin", "key"),
                "v": ("v_proj", "v_lin", "value"),
                "o": ("o_proj", "out_proj", "out_lin", "dense"),
            }
            for key, names in candidates.items():
                for n in names:
                    if hasattr(src, n):
                        m = getattr(src, n)
                        if m is not None and hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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("GLA UniversalAdapter fallback failed: %s", e)

        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

        with torch.no_grad():
            for key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target.weight)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target.weight.shape
                    and w.t().contiguous().shape == target.weight.shape
                ):
                    w = w.t().contiguous()

                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))

                if key in ("k", "v"):
                    src_out = w.shape[0]
                    tgt_out = target.weight.shape[0]
                    if src_out == tgt_out:
                        target.weight.data.copy_(w)
                    elif src_out > tgt_out and src_out % tgt_out == 0:
                        rep = src_out // tgt_out
                        target.weight.data.copy_(w.view(tgt_out, rep, -1).mean(dim=1))
                    elif src_out < tgt_out and tgt_out % src_out == 0:
                        rep = tgt_out // src_out
                        target.weight.data.copy_(w.repeat(rep, 1))
                    else:
                        target.weight.data.zero_()
                        target.weight.data[:out_d, :in_d].copy_(
                            w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                        )
                else:
                    target.weight.data.zero_()
                    target.weight.data[:out_d, :in_d].copy_(
                        w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                    )

                # GLA3-FIX-13: validate that the query projection is now fully
                # in the expected shape (all_head_size, hidden_size)
                if key == "q":
                    if target.weight.shape != (self.all_head_size, self.dim):
                        raise RuntimeError(
                            f"GLA copy_weights: q_proj shape {tuple(target.weight.shape)} "
                            f"!= expected ({self.all_head_size}, {self.dim}). "
                            f"num_heads={self.num_heads}, head_dim={self.head_dim}, dim={self.dim}."
                        )

                if b is not None and target.bias is not None:
                    mb = min(target.bias.numel(), b.numel())
                    target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))

                n = torch.linalg.norm(target.weight.data.float()).item()
                if n < 1e-4:
                    logger.warning("GLA: %s near-zero norm (%.2e)", key, n)

                target.freeze_base_weights()

            # Initialize z_proj from q
            w_q = extracted["q"][0]
            if w_q is not None:
                w_q_oriented = _orient(w_q, self.dim)
                out_d = min(self.z_proj.weight.shape[0], w_q_oriented.shape[0])
                in_d = min(self.z_proj.weight.shape[1], w_q_oriented.shape[1])
                self.z_proj.weight.data.zero_()
                self.z_proj.weight.data[:out_d, :in_d].copy_(
                    w_q_oriented[:out_d, :in_d].to(
                        self.z_proj.weight.device, self.z_proj.weight.dtype
                    )
                )
                self.z_proj.freeze_base_weights()

        try:
            for m in (
                self.q_proj,
                self.k_proj,
                self.v_proj,
                self.z_proj,
                self.out_proj,
            ):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

    # ------------------------------------------------------------------
    # Chunked linear-attention scan
    # ------------------------------------------------------------------
    def _gla_scan_chunked(self, q, k, v, lambdas, mask=None, reverse=False):
        B, H, L, d = q.shape
        device = q.device

        if HAS_TRITON_GLA and q.is_cuda and d <= MAX_REGISTER_DIM and mask is None and not reverse:
            try:
                return gla_scan_triton(q.float(), k.float(), v.float(), lambdas.float()).to(q.dtype)
            except Exception as e:
                logger.debug("GLA Triton scan failed, falling back: %s", e)

        if reverse:
            q = torch.flip(q, dims=[2])
            k = torch.flip(k, dims=[2])
            v = torch.flip(v, dims=[2])
            lambdas = torch.flip(lambdas, dims=[2])
            if mask is not None:
                mask = torch.flip(mask, dims=[1])

        if mask is not None:
            m = mask.float()
            m_kv = m.unsqueeze(1).unsqueeze(-1)
            k = k * m_kv
            v = v * m_kv

        q_f = q.float()
        k_f = k.float()
        v_f = v.float()
        lam = lambdas.float().clamp(min=1e-4, max=1.0 - 1e-4)

        chunk = self.chunk_size
        n_chunks = (L + chunk - 1) // chunk

        S = torch.zeros(B, H, d, d, device=device, dtype=torch.float32)
        Z = torch.zeros(B, H, d, device=device, dtype=torch.float32)
        outs = []

        for c in range(n_chunks):
            s = c * chunk
            e = min(s + chunk, L)
            q_c = q_f[:, :, s:e]
            k_c = k_f[:, :, s:e]
            v_c = v_f[:, :, s:e]
            lam_c = lam[:, :, s:e]

            log_lam = torch.log(lam_c)
            C = torch.cumsum(log_lam, dim=-1)
            C_i = C.unsqueeze(-1)
            C_j = C.unsqueeze(-2)
            D = torch.exp((C_i - C_j).clamp(min=-15.0, max=0.0)) * torch.tril(
                torch.ones(e - s, e - s, device=device)
            )

            lam_prefix = torch.exp(C)
            y_prev = torch.einsum(
                "bhl,bhld->bhld", lam_prefix, torch.einsum("bhlc,bhcd->bhld", q_c, S)
            )
            denom_prev = torch.einsum(
                "bhl,bhl->bhl", lam_prefix, torch.einsum("bhlc,bhc->bhl", q_c, Z)
            )

            qk = torch.matmul(q_c, k_c.transpose(-1, -2))
            y_cur = torch.matmul(qk * D, v_c)
            denom_cur = (qk * D).sum(dim=-1)

            y_chunk = y_prev + y_cur
            denom_chunk = (denom_prev + denom_cur).clamp(min=1e-4)

            lam_prod_chunk = torch.exp(C[:, :, -1]).view(B, H, 1, 1)
            C_end = C[:, :, -1:].unsqueeze(-1)
            C_t = C.unsqueeze(-1)
            weight_t = torch.exp((C_end - C_t).clamp(min=-15.0, max=0.0))
            k_w = k_c * weight_t
            S_new = lam_prod_chunk * S + torch.einsum("bhld,bhle->bhde", k_w, v_c)
            Z_new = lam_prod_chunk.squeeze(-1) * Z + k_w.sum(dim=2)

            S = S_new
            Z = Z_new

            y_chunk = y_chunk / denom_chunk.unsqueeze(-1)
            outs.append(y_chunk)

        y = torch.cat(outs, dim=2)
        if reverse:
            y = torch.flip(y, dims=[2])
        return y.to(q.dtype)

    # ------------------------------------------------------------------
    # Mask helpers (GLA3-FIX-10)
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_valid_mask(attention_mask, B, L, device):
        if attention_mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape == (B, L):
                return m
        elif attention_mask.dim() == 3:
            m = _to_bool(attention_mask)
            if m.shape[0] == B and m.shape[-1] == L:
                return m[:, -1, :]
        elif attention_mask.dim() == 4:
            m = _to_bool(attention_mask)
            # Shape (B, 1, N_q, L) or (B, H, N_q, L) → reduce over queries
            m = m.any(dim=-2)
            if m.dim() == 3:
                m = m.squeeze(1)
            if m.shape == (B, L):
                return m
        return None

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ):
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, L, D = hidden_states.shape

        # GLA3-FIX-05
        if self.dim != D:
            raise RuntimeError(
                f"GLA: hidden_states has last dim {D}, but module expects {self.dim}. "
                f"num_heads={self.num_heads}, head_dim={self.head_dim}."
            )

        H, d = self.num_heads, self.head_dim

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

        # ---- Projections ----
        q_raw = self.q_proj(hidden_states)
        k_raw = self.k_proj(kv_input)
        v_raw = self.v_proj(kv_input)

        # GLA3-FIX-07: verify projection outputs
        expected_q = (B, L, self.all_head_size)
        expected_kv = (B, L_kv, self.kv_head_size)
        if tuple(q_raw.shape) != expected_q:
            raise RuntimeError(f"GLA: q_proj output {tuple(q_raw.shape)} != expected {expected_q}")
        if tuple(k_raw.shape) != expected_kv:
            raise RuntimeError(f"GLA: k_proj output {tuple(k_raw.shape)} != expected {expected_kv}")

        q = q_raw.view(B, L, H, d).transpose(1, 2).contiguous()
        k = k_raw.view(B, L_kv, self.num_kv_heads, d).transpose(1, 2).contiguous()
        v = v_raw.view(B, L_kv, self.num_kv_heads, d).transpose(1, 2).contiguous()
        z = self.z_proj(hidden_states)

        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
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        if self.num_kv_heads < self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        # ELU+1 feature map
        q = F.elu(q) + 1.0
        k = F.elu(k) + 1.0

        # ---- Decay gate ----
        # GLA3-FIX-03 / GLA3-FIX-08: shape is guaranteed to be (B, L, num_heads)
        decay_out = self.decay_gate(hidden_states)  # (B, L, H)
        if decay_out.shape[-1] != self.num_heads:
            raise RuntimeError(
                f"GLA: decay_gate output {decay_out.shape[-1]} != num_heads {self.num_heads}"
            )
        decay_logits = decay_out.transpose(1, 2) + self.log_decay.to(decay_out.dtype).view(1, -1, 1)
        lambdas = torch.sigmoid(decay_logits)  # (B, H, L)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        mask_bool = self._normalize_valid_mask(eff_mask, B, L_kv, hidden_states.device)

        is_causal = bool(self.is_decoder and not is_cross and L > 1 and past_key_value is None)

        if is_causal:
            y = self._gla_scan_chunked(q, k, v, lambdas, mask=mask_bool, reverse=False)
        else:
            y_fwd = self._gla_scan_chunked(q, k, v, lambdas, mask=mask_bool, reverse=False)
            y_bwd = self._gla_scan_chunked(q, k, v, lambdas, mask=mask_bool, reverse=True)
            y = 0.5 * (y_fwd + y_bwd)

        y = self.state_norm(y)

        # GLA3-FIX-06: unified head_mask handling
        if head_mask is not None:
            hm = head_mask.to(y.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, H, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], H, 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            # 4D passes through
            if hm.shape[1] != H:
                raise RuntimeError(f"GLA: head_mask has {hm.shape[1]} heads, expected {H}")
            y = y * hm

        y = y.transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        context = torch.nan_to_num(y * F.silu(z), nan=0.0, posinf=1e4, neginf=-1e4)

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

        out = out * self.output_scale

        if mask_bool is not None and not is_cross:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        return out, None

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("q_proj", self.q_proj),
            ("k_proj", self.k_proj),
            ("v_proj", self.v_proj),
            ("z_proj", self.z_proj),
            ("out_proj", self.out_proj),
            ("decay_gate", self.decay_gate),
            ("state_norm", self.state_norm),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["output_scale"] = _sz(self.output_scale)
        out["total"] = sum(out.values())
        return out

Every 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.

var triton_kernel : str

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 storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("q_proj", self.q_proj),
        ("k_proj", self.k_proj),
        ("v_proj", self.v_proj),
        ("z_proj", self.z_proj),
        ("out_proj", self.out_proj),
        ("decay_gate", self.decay_gate),
        ("state_norm", self.state_norm),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["output_scale"] = _sz(self.output_scale)
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class GatedRecurrentMemory (d_model: int)

Expand source code

class GatedRecurrentMemory(nn.Module):
    """
    A light gated recurrent memory:
        g_t = sigmoid(W_g x_t + b_g)
        h_t = (1 - g_t) * h_{t-1} + g_t * tanh(W_x x_t + b_x)
        y_t = h_t

    Not a full DeltaNet but has the required gate structure and behaves as
    a proper recurrent memory (unlike the previous plain Linear layer).
    """

    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        self.gate_proj = nn.Linear(d_model, d_model, bias=True)
        self.value_proj = nn.Linear(d_model, d_model, bias=True)
        # Start gate slightly open so that memory doesn't dominate attention
        nn.init.constant_(self.gate_proj.bias, -1.0)
        nn.init.zeros_(self.value_proj.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, L, D)
        target_dtype = x.dtype
        gw = self.gate_proj.weight.to(target_dtype)
        gb = self.gate_proj.bias.to(target_dtype) if self.gate_proj.bias is not None else None
        vw = self.value_proj.weight.to(target_dtype)
        vb = self.value_proj.bias.to(target_dtype) if self.value_proj.bias is not None else None
        g = torch.sigmoid(F.linear(x, gw, gb))
        v = torch.tanh(F.linear(x, vw, vb))
        # Parallel scan over sequence (small L → simple python loop)
        B, L, D = x.shape
        h = torch.zeros(B, D, device=x.device, dtype=target_dtype)
        outs = []
        for t in range(L):
            h = (1.0 - g[:, t]) * h + g[:, t] * v[:, t]
            outs.append(h)
        return torch.stack(outs, dim=1).to(target_dtype)

A light gated recurrent memory: g_t = sigmoid(W_g x_t + b_g) h_t = (1 - g_t) * h_{t-1} + g_t * tanh(W_x x_t + b_x) y_t = h_t

Not a full DeltaNet but has the required gate structure and behaves as a proper recurrent memory (unlike the previous plain Linear layer).

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:
    # x: (B, L, D)
    target_dtype = x.dtype
    gw = self.gate_proj.weight.to(target_dtype)
    gb = self.gate_proj.bias.to(target_dtype) if self.gate_proj.bias is not None else None
    vw = self.value_proj.weight.to(target_dtype)
    vb = self.value_proj.bias.to(target_dtype) if self.value_proj.bias is not None else None
    g = torch.sigmoid(F.linear(x, gw, gb))
    v = torch.tanh(F.linear(x, vw, vb))
    # Parallel scan over sequence (small L → simple python loop)
    B, L, D = x.shape
    h = torch.zeros(B, D, device=x.device, dtype=target_dtype)
    outs = []
    for t in range(L):
        h = (1.0 - g[:, t]) * h + g[:, t] * v[:, t]
        outs.append(h)
    return torch.stack(outs, dim=1).to(target_dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class GatingModule (hidden_size, num_heads, head_dim, use_query, use_input)

Expand source code

class GatingModule(nn.Module):
    def __init__(self, hidden_size, num_heads, head_dim, use_query, use_input):
        super().__init__()
        self.query_gate = nn.Linear(head_dim, 1, bias=False) if use_query else None
        if use_input:
            self.input_gate = nn.Linear(hidden_size, num_heads, bias=True)
            nn.init.constant_(self.input_gate.bias, 1.8)
        else:
            self.input_gate = None

    def forward(self, hidden_states, Q):
        q_gate = torch.sigmoid(self.query_gate(Q)) if self.query_gate is not None else None
        in_gate = None
        if self.input_gate is not None:
            in_gate = torch.sigmoid(self.input_gate(hidden_states)).permute(0, 2, 1).unsqueeze(-1)
        return q_gate, in_gate

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, hidden_states, Q) ‑> Callable[..., typing.Any]

Expand source code

def forward(self, hidden_states, Q):
    q_gate = torch.sigmoid(self.query_gate(Q)) if self.query_gate is not None else None
    in_gate = None
    if self.input_gate is not None:
        in_gate = torch.sigmoid(self.input_gate(hidden_states)).permute(0, 2, 1).unsqueeze(-1)
    return q_gate, in_gate

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class GoldFinchAHNBlock (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
num_heads: int | None = None,
window_size: int = 64,
is_causal: bool = False,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
gate_init_bias: float = 1.8,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "goldfinch_ahn",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "goldfinch", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
class GoldFinchAHNBlock(BaseAttention):
    """
    GoldFinch-AHN: hybrid sliding-window attention + gated recurrent memory.

    Semantics of `window_size`: full window width (in tokens). Effective
    radius = `window_size // 2`, so `window_size = 64` → ±32 tokens.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        window_size: int = 64,
        is_causal: bool = False,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        gate_init_bias: float = 1.8,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # GF-FIX-05: None-safe
        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        # ---- hidden size ----
        if d_model is None:
            d_model = kwargs.get("d_model") or kwargs.get("hidden_size") or kwargs.get("dim")
            if d_model is None:
                if config is not None:
                    d_model = getattr(config, "hidden_size", getattr(config, "dim", 768))
                elif source_attn is not None:
                    out_layer = None
                    for attr in ("out_lin", "dense", "c_proj", "o_proj", "out_proj"):
                        if hasattr(source_attn, attr):
                            out_layer = getattr(source_attn, attr)
                            break
                    if out_layer is not None:
                        d_model = getattr(
                            out_layer,
                            "out_features",
                            getattr(out_layer, "in_features", 768),
                        )
                    else:
                        for attr in ("q_proj", "query", "q_lin"):
                            q_l = getattr(source_attn, attr, None)
                            if q_l is not None:
                                d_model = getattr(
                                    q_l,
                                    "in_features",
                                    getattr(q_l, "out_features", 768),
                                )
                                break
                if d_model is None:
                    d_model = 768

        if hasattr(d_model, "hidden_size"):
            d_model = d_model.hidden_size

        # ---- num_heads ----
        if num_heads is None:
            num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")
            if num_heads is None:
                if config is not None:
                    num_heads = getattr(
                        config, "num_attention_heads", getattr(config, "n_heads", 12)
                    )
                elif source_attn is not None:
                    num_heads = getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(source_attn, "n_heads", 12),
                    )
                else:
                    num_heads = 12

        self.dim = int(d_model)
        self.hidden_size = self.dim
        self.num_heads = max(1, int(num_heads))

        # ---- num_kv_heads ----
        if config is not None:
            self.num_kv_heads = getattr(
                config,
                "num_key_value_heads",
                getattr(config, "num_kv_heads", kwargs.get("num_kv_heads", self.num_heads)),
            )
        elif source_attn is not None:
            self.num_kv_heads = (
                getattr(source_attn, "num_key_value_heads", None)
                or getattr(source_attn, "num_kv_heads", None)
                or kwargs.get("num_kv_heads", self.num_heads)
            )
        else:
            self.num_kv_heads = kwargs.get("num_kv_heads", self.num_heads)
        self.num_kv_heads = max(1, min(int(self.num_kv_heads), self.num_heads))

        # ---- head_dim ----
        _head_dim = getattr(config, "head_dim", None) if config else None
        if _head_dim is None and source_attn is not None:
            _head_dim = getattr(source_attn, "head_dim", None)
        if _head_dim is None:
            _head_dim = self.dim // self.num_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 // self.num_kv_heads)

        # ---- window (GF-FIX-09) ----
        if config is not None:
            self.window_size = int(
                getattr(
                    config,
                    "window_size",
                    (
                        getattr(
                            comp_config,
                            "window_size",
                            kwargs.get("window_size", window_size),
                        )
                        if comp_config is not None
                        else window_size
                    ),
                )
            )
        elif comp_config is not None:
            self.window_size = int(
                getattr(comp_config, "window_size", kwargs.get("window_size", window_size))
            )
        else:
            self.window_size = int(kwargs.get("window_size", window_size))
        self.window_radius = max(1, self.window_size // 2)

        # ---- decoder detection ----
        model_type = str(getattr(config, "model_type", "") or "").lower() if config else ""
        is_dec = False
        if config is not None:
            is_dec = bool(
                getattr(config, "is_decoder", False) or getattr(config, "is_causal", False)
            )
        if any(
            model_type.startswith(p)
            for p in (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False) or is_causal:
            is_dec = True
        if source_attn is not None:
            if hasattr(source_attn, "c_attn") or getattr(source_attn, "is_decoder", False):
                is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # ---- LoRA config ----
        if config is not None:
            self.lora_rank = int(
                getattr(
                    config,
                    "lora_rank",
                    (
                        getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16))
                        if comp_config is not None
                        else kwargs.get("lora_rank", 16)
                    ),
                )
            )
            self.lora_alpha = float(
                getattr(
                    config,
                    "lora_alpha",
                    (
                        getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
                        if comp_config is not None
                        else kwargs.get("lora_alpha", 32.0)
                    ),
                )
            )
        elif comp_config is not None:
            self.lora_rank = int(getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)))
            self.lora_alpha = float(
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
            )
        else:
            self.lora_rank = int(kwargs.get("lora_rank", 16))
            self.lora_alpha = float(kwargs.get("lora_alpha", 32.0))

        self.freeze_base = self.lora_rank > 0
        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )

        # ---- projections ----
        self.W_q = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_k = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_v = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # GF-FIX-06: proper gated recurrent memory
        self.delta_mem = GatedRecurrentMemory(self.dim)

        # per-head attention / memory gate
        self.gate = nn.Linear(self.dim, self.num_heads)
        with torch.no_grad():
            self.gate.bias.data.fill_(gate_init_bias)
            nn.init.zeros_(self.gate.weight)

        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        dropout_p = 0.1
        if config is not None:
            dropout_p = float(
                getattr(
                    config,
                    "dropout",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                )
            )
        elif comp_config is not None:
            dropout_p = float(getattr(comp_config, "dropout", 0.1))
        self.dropout = nn.Dropout(dropout_p)

        # GF-FIX-11: single output scale, initialized to preserve pretrained signal
        self.output_scale = nn.Parameter(torch.ones(1))

        # Storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

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

    # ---- aliases ----
    @property
    def query(self):
        return self.W_q

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

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

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

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

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

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

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

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

    @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 freeze_base_weights(self):
        for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
            if hasattr(m, "freeze_base_weights"):
                m.freeze_base_weights()

    def unfreeze_base_weights(self):
        for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
            if hasattr(m, "unfreeze_base_weights"):
                m.unfreeze_base_weights()

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

        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        cls_name = getattr(src, "__class__", type(src)).__name__
        has_out = any(hasattr(src, a) for a in ("dense", "out_proj", "out_lin", "c_proj", "o_proj"))
        is_nested = False
        if (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
                "DebertaSelfAttention",
            )
            or hasattr(src, "query")
            and not has_out
            or hasattr(src, "self")
            and hasattr(src, "output")
        ):
            is_nested = True
        self.is_nested_encoder = bool(is_nested)

        extracted = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused
        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 = (
                    "Conv1D" in c_attn_mod.__class__.__name__
                    or "Conv1D" in base_c.__class__.__name__
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                )
                if is_conv:
                    c_w = c_w.t().contiguous()
                split = 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], c_b[:split] if c_b is not None else None)
                extracted["k"] = (
                    c_w[split : 2 * split],
                    c_b[split : 2 * split] if c_b is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :],
                    c_b[2 * split :] 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 = (
                        "Conv1D" in c_proj_mod.__class__.__name__
                        or "Conv1D" in base_o.__class__.__name__
                        or hasattr(c_proj_mod, "nf")
                        or hasattr(base_o, "nf")
                        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,
                    )

        # Nested (BERT / RoBERTa)
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                d = src.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat
        else:
            candidates = {
                "q": ("q_lin", "q_proj", "query", "W_q"),
                "k": ("k_lin", "k_proj", "key", "W_k"),
                "v": ("v_lin", "v_proj", "value", "W_v"),
                "o": ("out_lin", "out_proj", "dense", "o_proj"),
            }
            for key, names in candidates.items():
                for n in names:
                    if hasattr(src, n):
                        m = getattr(src, n)
                        if m is not None and hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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("Goldfinch 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:
            logger.warning(
                "Goldfinch copy_weights: missing %s; using initialized weights",
                missing_after,
            )

        target_map = {"q": self.W_q, "k": self.W_k, "v": self.W_v, "o": self.out_proj}

        with torch.no_grad():
            for key, target in target_map.items():
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target.weight)
                        if target.bias is not None:
                            nn.init.zeros_(target.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target.weight.shape
                    and w.t().contiguous().shape == target.weight.shape
                ):
                    w = w.t().contiguous()

                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))
                target.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                )

                if b is not None and target.bias is not None:
                    mb = min(target.bias.numel(), b.numel())
                    target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))
                elif target.bias is not None:
                    nn.init.zeros_(target.bias)

                # GF-FIX-02: warn, do not raise
                w_norm = float(torch.linalg.norm(target.weight.data.float()).item())
                if w_norm < 1e-4:
                    logger.warning(
                        "Goldfinch copy_weights: '%s' has near-zero norm (%.2e)",
                        key,
                        w_norm,
                    )

                if self.freeze_base:
                    target.freeze_base_weights()

        # Track compressed storage
        try:
            for m in (self.W_q, self.W_k, self.W_v, self.out_proj, self.delta_mem):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # GF-FIX-15: freeze once at the end, only if needed
        if self.freeze_base:
            self.freeze_base_weights()

    @classmethod
    def from_distilbert_pretrained(cls, old_attn, model_config=None, **kwargs):
        return cls(config=model_config, pretrained_attn=old_attn, **kwargs)

    # ------------------------------------------------------------------
    # Mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_valid_mask(attention_mask, B, L, device):
        if attention_mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape == (B, L):
                return m
        elif attention_mask.dim() == 4:
            m = _to_bool(attention_mask)
            m = m.any(dim=-2).squeeze(1)
            if m.shape == (B, L):
                return m

        try:
            m = get_token_valid_mask(attention_mask, B, L, device)
            if isinstance(m, torch.Tensor) and m.shape == (B, L):
                return m.bool()
        except Exception:
            pass
        return None

    # ------------------------------------------------------------------
    # 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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ):
        # GF-FIX-04: check parameter directly
        if encoder_hidden_states is not None:
            raise NotImplementedError("GoldFinchAHN does not support cross-attention.")

        x = hidden_states
        if not x.is_contiguous():
            x = x.contiguous()
        B, L, D = x.shape
        H, d = self.num_heads, self.head_dim
        device = x.device
        dtype_in = x.dtype

        # ---- QKV ----
        q = self.W_q(x).view(B, L, H, d).transpose(1, 2).contiguous()
        k = self.W_k(x).view(B, L, self.num_kv_heads, d).transpose(1, 2).contiguous()
        v = self.W_v(x).view(B, L, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # ---- RoPE ----
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)
        elif self.rotary_emb is not None:
            pos_ids = kwargs.get("position_ids")
            if pos_ids is None:
                pos_ids = torch.arange(L, device=device).unsqueeze(0).expand(B, -1)
            try:
                cos, sin = self.rotary_emb(k, seq_len=L)
            except TypeError:
                try:
                    cos, sin = self.rotary_emb(k, pos_ids)
                except Exception:
                    cos, sin = None, None
            if cos is not None and sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=pos_ids)

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

        # ---- Sliding-window attention ----
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * (d**-0.5)

        idx = torch.arange(L, device=device)
        diff = (idx.unsqueeze(1) - idx.unsqueeze(0)).abs()
        band = diff <= self.window_radius

        if self.is_causal:
            causal = idx.unsqueeze(1) >= idx.unsqueeze(0)
            valid_local = causal & band
        else:
            valid_local = band

        scores = scores.masked_fill(~valid_local.unsqueeze(0).unsqueeze(0), -1e4)

        valid_mask = self._normalize_valid_mask(attention_mask, B, L, device)
        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(dtype_in)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights)

        # GF-FIX-03 / GF-FIX-13: head_mask applied to per-head attention
        if head_mask is not None:
            hm = head_mask.to(attn_weights.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, H, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], H, 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            attn_weights = attn_weights * hm

        short_out = (
            torch.matmul(attn_weights, v)
            .transpose(1, 2)
            .contiguous()
            .view(B, L, self.all_head_size)
        )

        # ---- Recurrent memory ----
        long_out = self.delta_mem(x).to(dtype_in)  # (B, L, D)

        # ---- Gated fusion ----
        # gate: (B, L, H) → (B, L, 1) after mean
        gate_w = self.gate.weight.to(dtype_in)
        gate_b = self.gate.bias.to(dtype_in) if self.gate.bias is not None else None
        gate = (
            torch.sigmoid(F.linear(x, gate_w, gate_b)).mean(dim=-1, keepdim=True).to(dtype_in)
        )  # (B, L, 1)

        # Pad long_out to all_head_size if needed (should match since dim = all_head_size for MHA)
        if long_out.shape[-1] != self.all_head_size:
            if long_out.shape[-1] < self.all_head_size:
                long_out = F.pad(long_out, (0, self.all_head_size - long_out.shape[-1]))
            else:
                long_out = long_out[..., : self.all_head_size]

        fused = (gate * short_out.to(dtype_in) + (1.0 - gate) * long_out).to(dtype_in)

        # Zero-out padded positions
        if valid_mask is not None:
            fused = fused * valid_mask.unsqueeze(-1).to(fused.dtype)

        # ---- Output projection ----
        if self.is_nested_encoder:
            out = fused
        else:
            out = self.out_proj(fused)
            # GF-FIX-11: apply output_scale once
            out = (out * self.output_scale.to(out.dtype)).to(out.dtype)

        return out.to(dtype_in), (attn_weights if output_attentions else None)

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("W_q", self.W_q),
            ("W_k", self.W_k),
            ("W_v", self.W_v),
            ("out_proj", self.out_proj),
            ("delta_mem", self.delta_mem),
            ("gate", self.gate),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["output_scale"] = _sz(self.output_scale)
        out["total"] = sum(out.values())
        return out

GoldFinch-AHN: hybrid sliding-window attention + gated recurrent memory.

Semantics of window_size: full window width (in tokens). Effective radius = window_size // 2, so window_size = 64 → ±32 tokens.

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.

Static methods

def from_distilbert_pretrained(old_attn, model_config=None, **kwargs)

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.W_k

prop k_proj

Expand source code

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

prop key

Expand source code

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

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.W_q

prop q_proj

Expand source code

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

prop query

Expand source code

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

prop v_lin

Expand source code

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

prop v_proj

Expand source code

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

prop value

Expand source code

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

Methods

def freeze_base_weights(self)

Expand source code

def freeze_base_weights(self):
    for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
        if hasattr(m, "freeze_base_weights"):
            m.freeze_base_weights()

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("W_q", self.W_q),
        ("W_k", self.W_k),
        ("W_v", self.W_v),
        ("out_proj", self.out_proj),
        ("delta_mem", self.delta_mem),
        ("gate", self.gate),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["output_scale"] = _sz(self.output_scale)
    out["total"] = sum(out.values())
    return out

def unfreeze_base_weights(self)

Expand source code

def unfreeze_base_weights(self):
    for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
        if hasattr(m, "unfreeze_base_weights"):
            m.unfreeze_base_weights()

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class GraniteHybridBlock (config: Any | None = None,
rank: int | None = None,
dim: int | None = None,
num_heads: int | None = None,
num_experts: int = 4,
top_k: int = 2,
ratio: int = 2,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("granite_hybrid")
@ATTENTION_REGISTRY.register("granite_hybrid_block")
class GraniteHybridBlock(BaseAttention):
    r"""Granite Hybrid Attention-SSM Layer with MoE Feed-Forward and LoRA Adaptation."""

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        dim: int | None = None,
        num_heads: int | None = None,
        num_experts: int = 4,
        top_k: int = 2,
        ratio: int = 2,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        src_attn = pretrained_attn if pretrained_attn is not None else old_attn

        def _get(k: str, default: Any = None) -> Any:
            for obj in (
                config,
                comp_config,
                getattr(config, "compression_config", None),
            ):
                if obj is not None and hasattr(obj, k):
                    val = getattr(obj, k)
                    if val is not None:
                        return val
            return kwargs.get(k, default)

        if dim is None:
            dim = _get("hidden_size", _get("d_model", _get("n_embd", _get("dim", None))))
        if dim is None and src_attn is not None:
            dim = getattr(src_attn, "hidden_size", getattr(src_attn, "embed_dim", None))
            if dim is None:
                for cand in ("q_proj", "q_lin", "query", "c_attn"):
                    m = getattr(src_attn, cand, None)
                    if m is not None and hasattr(m, "in_features"):
                        dim = m.in_features
                        break
        dim = dim or kwargs.get("hidden_size", kwargs.get("d_model", kwargs.get("dim", 768)))

        if num_heads is None:
            num_heads = _get("num_attention_heads", _get("num_heads", None))
        if num_heads is None and src_attn is not None:
            num_heads = getattr(
                src_attn, "num_heads", getattr(src_attn, "num_attention_heads", None)
            )
        num_heads = num_heads or kwargs.get("num_attention_heads", kwargs.get("num_heads", 8))

        self.dim: int = int(dim)
        self.num_heads: int = int(num_heads)

        resolved_head_dim = _get("head_dim", None)
        if resolved_head_dim is None and src_attn is not None:
            resolved_head_dim = getattr(src_attn, "head_dim", None)
        if resolved_head_dim is None:
            resolved_head_dim = self.dim // max(1, self.num_heads)
        self.head_dim: int = int(resolved_head_dim)

        kv_heads = _get("num_key_value_heads", None)
        if kv_heads is None and src_attn is not None:
            kv_heads = getattr(
                src_attn,
                "num_key_value_heads",
                getattr(src_attn, "num_kv_heads", getattr(src_attn, "n_kv_heads", None)),
            )
            if (
                kv_heads is None
                and hasattr(src_attn, "k_proj")
                and hasattr(src_attn.k_proj, "out_features")
                and self.head_dim > 0
            ):
                kv_heads = src_attn.k_proj.out_features // self.head_dim
        self.num_kv_heads: int = int(kv_heads if kv_heads is not None else self.num_heads)
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.all_head_size: int = self.num_heads * self.head_dim
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))

        # Decoder detection
        model_type = str(_get("model_type", "")).lower()
        is_dec = _get(
            "is_decoder",
            _get("is_causal", kwargs.get("is_causal", kwargs.get("is_decoder", False))),
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if any(_mt.startswith(p) for p in _decoder_prefixes):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)

        # LoRA setup
        self.lora_rank: int = int(_get("lora_rank", kwargs.get("lora_rank", lora_rank)))
        self.lora_alpha: float = float(_get("lora_alpha", kwargs.get("lora_alpha", lora_alpha)))
        self.freeze_base: bool = self.lora_rank > 0

        self.mamba: Mamba2 = Mamba2(d_model=self.dim)

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

        num_exp = int(_get("num_experts", kwargs.get("num_experts", num_experts)))
        tk = int(_get("top_k", kwargs.get("top_k", top_k)))
        int(_get("ratio", kwargs.get("ratio", ratio)))
        self.moe: MixtureOfExperts = MixtureOfExperts(self.dim, self.dim * 4, num_exp, tk)
        self.alpha = nn.Parameter(torch.tensor(3.0))
        self.moe_scale = nn.Parameter(torch.zeros(1))
        self.is_nested_encoder = False
        self.output_scale = nn.Parameter(torch.ones(1))

        self.rotary_emb = getattr(src_attn, "rotary_emb", None) if src_attn is not None else None
        self.comp_config = comp_config

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

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

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        try:
            adapter = UniversalAdapter()
            for p in ["q", "k", "v", "o"]:
                w, b = adapter.get_weight_bias(old_attn, p)
                if w is not None:
                    extracted[p] = (w, b)
        except Exception as e:
            print(f"copy_weights failed: {e}")
            logger.warning(f"copy_weights failed: {e}")

        if hasattr(old_attn, "c_attn") and extracted["q"][0] is None:
            c_w, c_b = _unpack_weight_bias(old_attn.c_attn)
            if c_w is not None:
                if c_w.dim() == 2 and c_w.shape[1] == 3 * self.dim:
                    c_w = c_w.t().contiguous()
                split_size = c_w.size(0) // 3
                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(old_attn, "c_proj") and extracted["o"][0] is None:
                o_w, o_b = _unpack_weight_bias(old_attn.c_proj)
                if o_w is not None:
                    if o_w.dim() == 2 and o_w.shape[0] != self.dim and o_w.shape[1] == self.dim:
                        o_w = o_w.t().contiguous()
                    extracted["o"] = (o_w, o_b)

        elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
            s = old_attn.self
            for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]:
                extracted[proj_key] = _unpack_weight_bias(getattr(s, sub_name, None))
            if hasattr(old_attn, "output") and hasattr(old_attn.output, "dense"):
                extracted["o"] = _unpack_weight_bias(old_attn.output.dense)

        else:
            candidates = {
                "q": ["q_proj", "q_lin", "query"],
                "k": ["k_proj", "k_lin", "key"],
                "v": ["v_proj", "v_lin", "value"],
                "o": ["out_proj", "out_lin", "o_proj", "dense"],
            }
            for proj_key, names in candidates.items():
                if extracted[proj_key][0] is None:
                    for name in names:
                        if hasattr(old_attn, name):
                            cw, cb = _unpack_weight_bias(getattr(old_attn, name))
                            if cw is not None:
                                extracted[proj_key] = (cw, cb)
                                break

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }
        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is not None:
                    if w.shape == target_mod.weight.shape:
                        target_mod.weight.copy_(w)
                    elif w.t().contiguous().shape == target_mod.weight.shape:
                        target_mod.weight.copy_(w.t().contiguous())
                    elif (
                        w.shape[0] < target_mod.weight.shape[0]
                        and target_mod.weight.shape[0] % w.shape[0] == 0
                    ):
                        repeats = target_mod.weight.shape[0] // w.shape[0]
                        target_mod.weight.copy_(w.repeat(repeats, 1))
                    else:
                        min_0 = min(target_mod.weight.shape[0], w.shape[0])
                        min_1 = min(target_mod.weight.shape[1], w.shape[1])
                        target_mod.weight.zero_()
                        target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                    if b is not None and getattr(target_mod, "bias", None) is not None:
                        if b.shape == target_mod.bias.shape:
                            target_mod.bias.copy_(b)
                        elif (
                            b.shape[0] < target_mod.bias.shape[0]
                            and target_mod.bias.shape[0] % b.shape[0] == 0
                        ):
                            repeats = target_mod.bias.shape[0] // b.shape[0]
                            target_mod.bias.copy_(b.repeat(repeats))

        if self.freeze_base:
            for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
                mod.freeze_base_weights()

    def _attn_forward(
        self,
        q_in: torch.Tensor,
        k_in: torch.Tensor,
        v_in: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L_q, _ = q_in.shape
        L_k = k_in.shape[1]
        H = self.num_heads
        H_kv = self.num_kv_heads
        Dh = self.head_dim

        q = self.q_proj(q_in).view(B, L_q, H, Dh).transpose(1, 2).contiguous()
        k = self.k_proj(k_in).view(B, L_k, H_kv, Dh).transpose(1, 2).contiguous()
        v = self.v_proj(v_in).view(B, L_k, H_kv, Dh).transpose(1, 2).contiguous()

        # RoPE application
        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        pos_ids = kwargs.get("position_ids")
        if pos_emb is not None:
            if isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
                cos, sin = pos_emb[0], pos_emb[1]
                q, k = apply_rotary_pos_emb(q, k, cos, sin, pos_ids)
        elif self.rotary_emb is not None:
            rotary_mod = self.rotary_emb
            if pos_ids is None:
                pos_ids = torch.arange(L_q, device=q_in.device).unsqueeze(0).expand(B, -1)
            try:
                cos, sin = rotary_mod(q, seq_len=L_k)
            except Exception:
                try:
                    cos, sin = rotary_mod(q, pos_ids)
                except Exception:
                    try:
                        cos, sin = rotary_mod(q)
                    except Exception:
                        cos, sin = None, None
            if cos is not None and sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=pos_ids)

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

        scale = 1.0 / math.sqrt(Dh)
        scores = torch.matmul(q, k.transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and L_q == L_k:
            causal_mask = torch.triu(torch.full((L_q, L_k), -1e4, device=q_in.device), diagonal=1)
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

        if attention_mask is not None:
            if attention_mask.dim() == 2:
                pad_mask = (attention_mask == 0).view(B, 1, 1, L_k)
                scores = scores.masked_fill(pad_mask, -10000.0)
            elif attention_mask.dim() == 3:
                scores = scores + attention_mask.unsqueeze(1)
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        context = torch.matmul(attn_weights, v)
        context = context.transpose(1, 2).contiguous().view(B, L_q, self.all_head_size)
        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_proj(context)
        return out, attn_weights

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if encoder_hidden_states is not None:
            return hidden_states, None

        while isinstance(hidden_states, (tuple, list)):
            hidden_states = hidden_states[0]

        attn_out, attn_weights = self._attn_forward(
            hidden_states,
            hidden_states,
            hidden_states,
            attention_mask=attention_mask,
            position_embeddings=position_embeddings,
            **kwargs,
        )

        mamba_out = self.mamba(hidden_states)
        gate = torch.sigmoid(self.alpha)
        x = gate * attn_out + (1.0 - gate) * mamba_out
        x = x + self.moe_scale * self.moe(x)

        if x.shape[-1] != hidden_states.shape[-1]:
            if x.shape[-1] < hidden_states.shape[-1]:
                x = F.pad(x, (0, hidden_states.shape[-1] - x.shape[-1]))
            else:
                x = x[..., : hidden_states.shape[-1]]

        return x, (attn_weights if output_attentions else None)

Granite Hybrid Attention-SSM Layer with MoE Feed-Forward and LoRA Adaptation.

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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class GriffinSSDBlock (config: Any | None = None,
rank: int | None = 64,
window_size: int = 64,
d_model: int | Any | None = None,
num_heads: int | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("griffin_ssd")
@ATTENTION_REGISTRY.register("griffin_ssd_block")
class GriffinSSDBlock(BaseAttention):
    r"""Griffin Hybrid Block combining local window attention with global SSD linear attention."""

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        window_size: int = 64,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        def _get(k: str, default: Any = None) -> Any:
            for obj in (
                comp_config,
                config,
                getattr(config, "compression_config", None),
            ):
                if obj is not None:
                    if isinstance(obj, dict) and k in obj and obj[k] is not None:
                        return obj[k]
                    elif hasattr(obj, k):
                        val = getattr(obj, k)
                        if val is not None:
                            return val
            return kwargs.get(k, default)

        # Resolve d_model
        if d_model is None:
            d_model = _get("hidden_size", _get("d_model", _get("n_embd", _get("dim", None))))
        if d_model is None and old_attn is not None:
            d_model = (
                getattr(old_attn, "hidden_size", None)
                or getattr(old_attn, "embed_dim", None)
                or getattr(old_attn, "all_head_size", None)
            )
            if d_model is None:
                for cand in ("q_proj", "q_lin", "query", "c_attn"):
                    m = getattr(old_attn, cand, None)
                    if m is not None and hasattr(m, "in_features"):
                        d_model = m.in_features
                        break
        if d_model is None:
            d_model = kwargs.get("hidden_size", kwargs.get("d_model", kwargs.get("dim", 768)))

        # Resolve num_heads
        if num_heads is None:
            num_heads = (
                _get("num_attention_heads", None)
                or _get("num_heads", None)
                or _get("n_heads", None)
                or _get("n_head", None)
            )
        if num_heads is None and old_attn is not None:
            num_heads = getattr(
                old_attn, "num_heads", getattr(old_attn, "num_attention_heads", None)
            )
        if num_heads is None:
            num_heads = kwargs.get("num_attention_heads", kwargs.get("num_heads", 12))

        self.dim: int = int(d_model)
        self.num_heads: int = int(num_heads)

        resolved_head_dim = _get("head_dim", None)
        if resolved_head_dim is None and old_attn is not None:
            resolved_head_dim = getattr(old_attn, "head_dim", None)
            if resolved_head_dim is None:
                for q_cand in ("q_proj", "q_lin", "query", "W_q"):
                    m = getattr(old_attn, q_cand, None)
                    if m is not None and hasattr(m, "out_features"):
                        resolved_head_dim = m.out_features // max(1, self.num_heads)
                        break
        if resolved_head_dim is None:
            resolved_head_dim = self.dim // max(1, self.num_heads)
        self.head_dim: int = int(resolved_head_dim)

        kv_heads = _get("num_key_value_heads", None)
        if kv_heads is None and old_attn is not None:
            kv_heads = getattr(
                old_attn,
                "num_key_value_heads",
                getattr(old_attn, "num_kv_heads", getattr(old_attn, "n_kv_heads", None)),
            )
            if (
                kv_heads is None
                and hasattr(old_attn, "k_proj")
                and hasattr(old_attn.k_proj, "out_features")
                and self.head_dim > 0
            ):
                kv_heads = old_attn.k_proj.out_features // self.head_dim
        self.num_kv_heads: int = int(kv_heads if kv_heads is not None else self.num_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
        self.window_size: int = int(_get("window_size", kwargs.get("window_size", window_size)))

        # Decoder / causal detection
        model_type = str(_get("model_type", "")).lower()
        is_dec = _get(
            "is_decoder",
            _get("is_causal", kwargs.get("is_causal", kwargs.get("is_decoder", False))),
        )
        if (
            any(
                model_type.startswith(pfx)
                for pfx in (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "bloom",
                    "opt",
                    "falcon",
                    "codegen",
                )
            )
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)

        # Nested encoder detection
        is_nested = False
        if old_attn is not None:
            src_cls = old_attn.__class__.__name__
            if (
                src_cls
                in [
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                    "ElectraSelfAttention",
                    "DebertaSelfAttention",
                ]
                or (src_cls.endswith("SelfAttention") and "distil" not in src_cls.lower())
                or (hasattr(old_attn, "self") and hasattr(old_attn, "output"))
            ):
                is_nested = True
        if (
            not is_nested
            and model_type in ["roberta", "bert", "xlm-roberta"]
            and not self.is_decoder
        ):
            is_nested = True
        self.is_nested_encoder: bool = bool(is_nested)

        # LoRA setup
        self.lora_rank: int = int(_get("lora_rank", kwargs.get("lora_rank", lora_rank)))
        self.lora_alpha: float = float(_get("lora_alpha", kwargs.get("lora_alpha", lora_alpha)))
        self.freeze_base: bool = self.lora_rank > 0

        # Multi-head projections
        self.W_q = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_k = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_v = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.ssd_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Learnable decay gate for global SSD branch (0.0 -> sigmoid(0) = 0.5)
        self.decay_param = nn.Parameter(torch.zeros(self.num_heads, 1, 1))

        self.rotary_emb = getattr(old_attn, "rotary_emb", None) if old_attn is not None else None
        self.comp_config = comp_config

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

    def copy_weights(self, old_attn: Any) -> None:
        """Transfer pretrained projection weights across multiple model families."""
        if old_attn is None:
            return

        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
        try:
            adapter = UniversalAdapter()
            for proj in ["q", "k", "v", "o"]:
                w, b = adapter.get_weight_bias(old_attn, proj)
                if w is not None:
                    extracted[proj] = (w, b)
        except Exception:
            pass

        # 2. GPT-2 Fused c_attn
        if hasattr(old_attn, "c_attn") and extracted["q"][0] is None:
            c_w, c_b = _unpack_weight_bias(old_attn.c_attn)
            if c_w is not None:
                if c_w.dim() == 2 and c_w.shape[1] == 3 * self.dim:
                    c_w = c_w.t().contiguous()
                split_size = c_w.size(0) // 3
                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(old_attn, "c_proj") and extracted["o"][0] is None:
                o_w, o_b = _unpack_weight_bias(old_attn.c_proj)
                if o_w is not None:
                    if o_w.dim() == 2 and o_w.shape[0] != self.dim and o_w.shape[1] == self.dim:
                        o_w = o_w.t().contiguous()
                    extracted["o"] = (o_w, o_b)

        # 3. RoBERTa / BERT Nested
        elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
            self.is_nested_encoder = True
            s = old_attn.self
            for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]:
                mod = getattr(s, sub_name, None)
                extracted[proj_key] = _unpack_weight_bias(mod)
            if hasattr(old_attn, "output") and hasattr(old_attn.output, "dense"):
                extracted["o"] = _unpack_weight_bias(old_attn.output.dense)

        # 4. Flat projections (DistilBERT, Qwen, LLaMA)
        else:
            candidates = {
                "q": ["q_proj", "q_lin", "query"],
                "k": ["k_proj", "k_lin", "key"],
                "v": ["v_proj", "v_lin", "value"],
                "o": ["out_proj", "out_lin", "o_proj", "dense"],
            }
            for proj_key, names in candidates.items():
                if extracted[proj_key][0] is None:
                    for name in names:
                        if hasattr(old_attn, name):
                            mod = getattr(old_attn, name)
                            cw, cb = _unpack_weight_bias(mod)
                            if cw is not None:
                                extracted[proj_key] = (cw, cb)
                                break

        # Copy into modules
        target_map = {
            "q": self.W_q,
            "k": self.W_k,
            "v": self.W_v,
            "o": self.out_proj,
        }
        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is not None:
                    if w.shape == target_mod.weight.shape:
                        target_mod.weight.copy_(w)
                    elif w.t().contiguous().shape == target_mod.weight.shape:
                        target_mod.weight.copy_(w.t().contiguous())
                    elif (
                        w.shape[0] < target_mod.weight.shape[0]
                        and target_mod.weight.shape[0] % w.shape[0] == 0
                    ):
                        repeats = target_mod.weight.shape[0] // w.shape[0]
                        target_mod.weight.copy_(w.repeat(repeats, 1))
                    else:
                        min_0 = min(target_mod.weight.shape[0], w.shape[0])
                        min_1 = min(target_mod.weight.shape[1], w.shape[1])
                        target_mod.weight.zero_()
                        target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                    if b is not None and getattr(target_mod, "bias", None) is not None:
                        if b.shape == target_mod.bias.shape:
                            target_mod.bias.copy_(b)
                        elif (
                            b.shape[0] < target_mod.bias.shape[0]
                            and target_mod.bias.shape[0] % b.shape[0] == 0
                        ):
                            repeats = target_mod.bias.shape[0] // b.shape[0]
                            target_mod.bias.copy_(b.repeat(repeats))

            # Initialize ssd_proj from q
            if extracted["q"][0] is not None:
                w_q = extracted["q"][0]
                if w_q.shape == self.ssd_proj.weight.shape:
                    self.ssd_proj.weight.copy_(w_q)
                elif w_q.t().contiguous().shape == self.ssd_proj.weight.shape:
                    self.ssd_proj.weight.copy_(w_q.t().contiguous())

        if self.freeze_base:
            for mod in [self.W_q, self.W_k, self.W_v, self.ssd_proj, self.out_proj]:
                mod.freeze_base_weights()

    def _ssd_scan_parallel(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        decay: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """State-Space Duality (SSD) matrix scan supporting forward causal and bidirectional mode."""
        B, H, L, d = q.shape
        target_dtype = q.dtype

        decay_typed = decay.to(target_dtype)
        log_decay = torch.log(decay_typed.clamp(min=1e-4, max=1.0 - 1e-4))  # [1, H, 1, 1]
        pos = torch.arange(L, device=q.device, dtype=target_dtype)
        pos_diff = pos.unsqueeze(-1) - pos.unsqueeze(-2)  # [L, L]

        attn_scores = (q @ k.transpose(-1, -2).contiguous()).to(target_dtype)  # [B, H, L, L]

        if mask is not None:
            mask_typed = (
                mask.to(target_dtype) if mask.dim() == 2 else mask[:, 0, 0, :].to(target_dtype)
            )
            mask_2d = mask_typed.unsqueeze(1).unsqueeze(-1) * mask_typed.unsqueeze(1).unsqueeze(-2)
        else:
            mask_2d = None

        if self.is_decoder:
            # Forward causal decay matrix
            tril_mask = torch.tril(torch.ones(L, L, device=q.device, dtype=target_dtype))
            D_causal = (
                torch.exp(
                    (pos_diff.unsqueeze(0).unsqueeze(0) * log_decay).clamp(min=-40.0, max=0.0)
                ).to(target_dtype)
                * tril_mask
            )
            if mask_2d is not None:
                D_causal = D_causal * mask_2d
            D_causal = D_causal.to(target_dtype)
            v_typed = v.to(target_dtype)
            y_fwd = (attn_scores * D_causal) @ v_typed  # [B, H, L, d]
            return y_fwd.to(target_dtype), None
        else:
            # Symmetric distance decay for bidirectional encoders: D_sym(i, j) = exp(|i - j| * log_decay)
            dist = torch.abs(pos_diff).unsqueeze(0).unsqueeze(0)
            D_sym = torch.exp((dist * log_decay).clamp(min=-40.0, max=0.0)).to(target_dtype)
            if mask_2d is not None:
                D_sym = D_sym * mask_2d
            D_sym = D_sym.to(target_dtype)
            v_typed = v.to(target_dtype)
            y_sym = (attn_scores * D_sym) @ v_typed  # [B, H, L, d]
            return y_sym.to(target_dtype), None

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Forward pass for Griffin SSD attention."""
        if encoder_hidden_states is not None:
            return hidden_states, None

        mask = kwargs.get("mask")
        if attention_mask is None and mask is not None:
            attention_mask = mask

        while isinstance(hidden_states, (tuple, list)):
            hidden_states = hidden_states[0]

        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        target_dtype = hidden_states.dtype
        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim

        # 1. Multi-head Projections
        q = self.W_q(hidden_states).view(B, L, H, d).transpose(1, 2).contiguous().to(target_dtype)
        k = (
            self.W_k(hidden_states)
            .view(B, L, self.num_kv_heads, d)
            .transpose(1, 2)
            .contiguous()
            .to(target_dtype)
        )
        v = (
            self.W_v(hidden_states)
            .view(B, L, self.num_kv_heads, d)
            .transpose(1, 2)
            .contiguous()
            .to(target_dtype)
        )
        q_ssd = (
            self.ssd_proj(hidden_states)
            .view(B, L, H, d)
            .transpose(1, 2)
            .contiguous()
            .to(target_dtype)
        )

        # Apply RoPE if provided or probe from rotary_emb
        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        pos_ids = kwargs.get("position_ids")
        if pos_emb is not None:
            if isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
                cos, sin = pos_emb[0], pos_emb[1]
            else:
                cos, sin = pos_emb, None
            if cos is not None and sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=pos_ids)
        elif self.rotary_emb is not None:
            rotary_mod = self.rotary_emb
            if pos_ids is None:
                pos_ids = torch.arange(L, device=hidden_states.device).unsqueeze(0).expand(B, -1)
            cos, sin = None, None
            try:
                cos, sin = rotary_mod(q, seq_len=L)
            except Exception:
                try:
                    cos, sin = rotary_mod(q, pos_ids)
                except Exception:
                    try:
                        cos, sin = rotary_mod(q)
                    except Exception:
                        cos, sin = None, None
            if cos is not None and sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=pos_ids)

        # GQA repeat
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = repeat_kv(k, repeats)
            v = repeat_kv(v, repeats)

        # 2. Local Sliding Window Attention
        scores = torch.matmul(q, k.transpose(-1, -2).contiguous()) / math.sqrt(d)  # [B, H, L, L]
        idx = torch.arange(L, device=hidden_states.device)
        dist = torch.abs(idx.unsqueeze(1) - idx.unsqueeze(0))
        window_mask = dist > self.window_size

        mask_val = -10000.0
        scores = scores.masked_fill(window_mask.unsqueeze(0).unsqueeze(0), mask_val)

        if self.is_decoder:
            causal_mask = torch.triu(
                torch.ones(L, L, device=hidden_states.device, dtype=torch.bool),
                diagonal=1,
            )
            scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), mask_val)

        mask_bool = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                mask_bool = attention_mask.bool()
                pad_mask = (attention_mask == 0).view(B, 1, 1, L)
                scores = scores.masked_fill(pad_mask, mask_val)
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask.to(scores.dtype)

        local_attn = F.softmax(scores, dim=-1).to(target_dtype)
        local_attn = torch.nan_to_num(local_attn, nan=0.0)
        local_out = torch.matmul(local_attn, v).to(target_dtype)  # [B, H, L, d]

        # 3. Global SSD Recurrence
        q_norm = F.normalize(q_ssd, p=2, dim=-1).to(target_dtype)
        k_norm = (F.normalize(k, p=2, dim=-1) / (d**0.5)).to(target_dtype)
        decay = torch.sigmoid(self.decay_param.to(target_dtype)).view(1, H, 1, 1)

        ssd_fwd, ssd_bwd = self._ssd_scan_parallel(q_norm, k_norm, v, decay, mask_bool)

        if self.is_decoder or ssd_bwd is None:
            global_out = ssd_fwd.to(target_dtype)
        else:
            global_out = (0.5 * (ssd_fwd + ssd_bwd)).to(target_dtype)

        # 4. Fused Output
        fused = (
            (local_out + global_out)
            .transpose(1, 2)
            .contiguous()
            .view(B, L, self.all_head_size)
            .to(target_dtype)
        )
        if self.is_nested_encoder:
            out = fused
        else:
            out = self.out_proj(fused).to(target_dtype)

        if out.shape[-1] != hidden_states.shape[-1]:
            if out.shape[-1] < hidden_states.shape[-1]:
                out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
            else:
                out = out[..., : hidden_states.shape[-1]]

        if attention_mask is not None and mask_bool is not None:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        return out.to(target_dtype), None

Griffin Hybrid Block combining local window attention with global SSD linear attention.

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.

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Transfer pretrained projection weights across multiple model families."""
    if old_attn is None:
        return

    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
    try:
        adapter = UniversalAdapter()
        for proj in ["q", "k", "v", "o"]:
            w, b = adapter.get_weight_bias(old_attn, proj)
            if w is not None:
                extracted[proj] = (w, b)
    except Exception:
        pass

    # 2. GPT-2 Fused c_attn
    if hasattr(old_attn, "c_attn") and extracted["q"][0] is None:
        c_w, c_b = _unpack_weight_bias(old_attn.c_attn)
        if c_w is not None:
            if c_w.dim() == 2 and c_w.shape[1] == 3 * self.dim:
                c_w = c_w.t().contiguous()
            split_size = c_w.size(0) // 3
            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(old_attn, "c_proj") and extracted["o"][0] is None:
            o_w, o_b = _unpack_weight_bias(old_attn.c_proj)
            if o_w is not None:
                if o_w.dim() == 2 and o_w.shape[0] != self.dim and o_w.shape[1] == self.dim:
                    o_w = o_w.t().contiguous()
                extracted["o"] = (o_w, o_b)

    # 3. RoBERTa / BERT Nested
    elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
        self.is_nested_encoder = True
        s = old_attn.self
        for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]:
            mod = getattr(s, sub_name, None)
            extracted[proj_key] = _unpack_weight_bias(mod)
        if hasattr(old_attn, "output") and hasattr(old_attn.output, "dense"):
            extracted["o"] = _unpack_weight_bias(old_attn.output.dense)

    # 4. Flat projections (DistilBERT, Qwen, LLaMA)
    else:
        candidates = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["out_proj", "out_lin", "o_proj", "dense"],
        }
        for proj_key, names in candidates.items():
            if extracted[proj_key][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        mod = getattr(old_attn, name)
                        cw, cb = _unpack_weight_bias(mod)
                        if cw is not None:
                            extracted[proj_key] = (cw, cb)
                            break

    # Copy into modules
    target_map = {
        "q": self.W_q,
        "k": self.W_k,
        "v": self.W_v,
        "o": self.out_proj,
    }
    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is not None:
                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and getattr(target_mod, "bias", None) is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))

        # Initialize ssd_proj from q
        if extracted["q"][0] is not None:
            w_q = extracted["q"][0]
            if w_q.shape == self.ssd_proj.weight.shape:
                self.ssd_proj.weight.copy_(w_q)
            elif w_q.t().contiguous().shape == self.ssd_proj.weight.shape:
                self.ssd_proj.weight.copy_(w_q.t().contiguous())

    if self.freeze_base:
        for mod in [self.W_q, self.W_k, self.W_v, self.ssd_proj, self.out_proj]:
            mod.freeze_base_weights()

Transfer pretrained projection weights across multiple model families.

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Forward pass for Griffin SSD attention."""
    if encoder_hidden_states is not None:
        return hidden_states, None

    mask = kwargs.get("mask")
    if attention_mask is None and mask is not None:
        attention_mask = mask

    while isinstance(hidden_states, (tuple, list)):
        hidden_states = hidden_states[0]

    if not hidden_states.is_contiguous():
        hidden_states = hidden_states.contiguous()

    target_dtype = hidden_states.dtype
    B, L, D = hidden_states.shape
    H, d = self.num_heads, self.head_dim

    # 1. Multi-head Projections
    q = self.W_q(hidden_states).view(B, L, H, d).transpose(1, 2).contiguous().to(target_dtype)
    k = (
        self.W_k(hidden_states)
        .view(B, L, self.num_kv_heads, d)
        .transpose(1, 2)
        .contiguous()
        .to(target_dtype)
    )
    v = (
        self.W_v(hidden_states)
        .view(B, L, self.num_kv_heads, d)
        .transpose(1, 2)
        .contiguous()
        .to(target_dtype)
    )
    q_ssd = (
        self.ssd_proj(hidden_states)
        .view(B, L, H, d)
        .transpose(1, 2)
        .contiguous()
        .to(target_dtype)
    )

    # Apply RoPE if provided or probe from rotary_emb
    pos_emb = (
        position_embeddings
        if position_embeddings is not None
        else kwargs.get("position_embeddings")
    )
    pos_ids = kwargs.get("position_ids")
    if pos_emb is not None:
        if isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
        else:
            cos, sin = pos_emb, None
        if cos is not None and sin is not None:
            q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=pos_ids)
    elif self.rotary_emb is not None:
        rotary_mod = self.rotary_emb
        if pos_ids is None:
            pos_ids = torch.arange(L, device=hidden_states.device).unsqueeze(0).expand(B, -1)
        cos, sin = None, None
        try:
            cos, sin = rotary_mod(q, seq_len=L)
        except Exception:
            try:
                cos, sin = rotary_mod(q, pos_ids)
            except Exception:
                try:
                    cos, sin = rotary_mod(q)
                except Exception:
                    cos, sin = None, None
        if cos is not None and sin is not None:
            q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=pos_ids)

    # GQA repeat
    if self.num_heads != self.num_kv_heads:
        repeats = self.num_heads // self.num_kv_heads
        k = repeat_kv(k, repeats)
        v = repeat_kv(v, repeats)

    # 2. Local Sliding Window Attention
    scores = torch.matmul(q, k.transpose(-1, -2).contiguous()) / math.sqrt(d)  # [B, H, L, L]
    idx = torch.arange(L, device=hidden_states.device)
    dist = torch.abs(idx.unsqueeze(1) - idx.unsqueeze(0))
    window_mask = dist > self.window_size

    mask_val = -10000.0
    scores = scores.masked_fill(window_mask.unsqueeze(0).unsqueeze(0), mask_val)

    if self.is_decoder:
        causal_mask = torch.triu(
            torch.ones(L, L, device=hidden_states.device, dtype=torch.bool),
            diagonal=1,
        )
        scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), mask_val)

    mask_bool = None
    if attention_mask is not None:
        if attention_mask.dim() == 2:
            mask_bool = attention_mask.bool()
            pad_mask = (attention_mask == 0).view(B, 1, 1, L)
            scores = scores.masked_fill(pad_mask, mask_val)
        elif attention_mask.dim() == 4:
            scores = scores + attention_mask.to(scores.dtype)

    local_attn = F.softmax(scores, dim=-1).to(target_dtype)
    local_attn = torch.nan_to_num(local_attn, nan=0.0)
    local_out = torch.matmul(local_attn, v).to(target_dtype)  # [B, H, L, d]

    # 3. Global SSD Recurrence
    q_norm = F.normalize(q_ssd, p=2, dim=-1).to(target_dtype)
    k_norm = (F.normalize(k, p=2, dim=-1) / (d**0.5)).to(target_dtype)
    decay = torch.sigmoid(self.decay_param.to(target_dtype)).view(1, H, 1, 1)

    ssd_fwd, ssd_bwd = self._ssd_scan_parallel(q_norm, k_norm, v, decay, mask_bool)

    if self.is_decoder or ssd_bwd is None:
        global_out = ssd_fwd.to(target_dtype)
    else:
        global_out = (0.5 * (ssd_fwd + ssd_bwd)).to(target_dtype)

    # 4. Fused Output
    fused = (
        (local_out + global_out)
        .transpose(1, 2)
        .contiguous()
        .view(B, L, self.all_head_size)
        .to(target_dtype)
    )
    if self.is_nested_encoder:
        out = fused
    else:
        out = self.out_proj(fused).to(target_dtype)

    if out.shape[-1] != hidden_states.shape[-1]:
        if out.shape[-1] < hidden_states.shape[-1]:
            out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
        else:
            out = out[..., : hidden_states.shape[-1]]

    if attention_mask is not None and mask_bool is not None:
        out = out * mask_bool.unsqueeze(-1).to(out.dtype)

    return out.to(target_dtype), None

Forward pass for Griffin SSD attention.

Inherited members

• BaseAttention:
  • extra_repr

class GroupedQueryAttention (config: Any | None = None,
rank: int | None = None,
dim: int | None = None,
num_heads: int | None = None,
num_kv_heads: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
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(
    "gqa",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
    supports_rope=True,
)
@ATTENTION_REGISTRY.register(
    "gqa_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
    supports_rope=True,
)
@ATTENTION_REGISTRY.register(
    "grouped_query_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
    supports_rope=True,
)
class GroupedQueryAttention(BaseAttention):
    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,
        dim: int | None = None,
        num_heads: int | None = None,
        num_kv_heads: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        src_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # ---- hidden_size ----
        if dim is None and config is not None:
            dim = (
                _cfg_get(config, "hidden_size", None)
                or _cfg_get(config, "d_model", None)
                or _cfg_get(config, "n_embd", None)
                or _cfg_get(config, "dim", None)
            )
        if dim is None and src_attn is not None:
            dim = getattr(src_attn, "hidden_size", getattr(src_attn, "embed_dim", None))
            if dim is None:
                for cand in ("q_proj", "q_lin", "query", "c_attn"):
                    m = getattr(src_attn, cand, None)
                    if m is not None:
                        w, _ = _unpack_weight_bias(m)
                        if w is not None and w.dim() == 2:
                            dim = w.shape[1]
                            break
        self.hidden_size = int(dim or kwargs.get("hidden_size", 768))

        # ---- num_heads ----
        if num_heads is None and config is not None:
            num_heads = (
                _cfg_get(config, "num_attention_heads", None)
                or _cfg_get(config, "num_heads", None)
                or _cfg_get(config, "n_head", None)
            )
        if num_heads is None and src_attn is not None:
            num_heads = getattr(
                src_attn,
                "num_heads",
                getattr(src_attn, "num_attention_heads", getattr(src_attn, "n_head", None)),
            )
        self.num_heads = int(
            num_heads or kwargs.get("num_attention_heads", kwargs.get("num_heads", 8))
        )

        # ---- head_dim ----
        self.head_dim = int(
            kwargs.get("head_dim", getattr(config, "head_dim", None) if config else None)
            or (self.hidden_size // max(1, self.num_heads))
        )

        # ==================================================================
        # GQA-FIX-16: derive num_kv_heads from `rank` when not explicitly given
        # ==================================================================
        # Priority:
        #   1. explicit `num_kv_heads` argument
        #   2. config's num_key_value_heads
        #   3. `rank` argument (1 <= rank <= num_heads)
        #   4. kwargs' num_kv_heads
        #   5. default = num_heads // 4
        # ==================================================================
        resolved_kv_heads = None
        source = "default"

        if num_kv_heads is not None:
            resolved_kv_heads = int(num_kv_heads)
            source = "explicit_arg"
        elif config is not None:
            cfg_kv = _cfg_get(config, "num_key_value_heads", None) or _cfg_get(
                config, "num_kv_heads", None
            )
            if cfg_kv is not None:
                resolved_kv_heads = int(cfg_kv)
                source = "config"
        if resolved_kv_heads is None and src_attn is not None:
            src_kv = getattr(
                src_attn, "num_key_value_heads", getattr(src_attn, "num_kv_heads", None)
            )
            if src_kv is not None:
                resolved_kv_heads = int(src_kv)
                source = "src_attn"

        # GQA-FIX-16: if `rank` is given and looks like a head count (1..num_heads),
        # use it as num_kv_heads. This makes the rank sweep actually produce
        # different compression levels.
        if rank is not None and 1 <= int(rank) <= self.num_heads:
            resolved_kv_heads = int(rank)
            source = "rank"

        if resolved_kv_heads is None:
            kwargs_kv = kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads"))
            if kwargs_kv is not None:
                resolved_kv_heads = int(kwargs_kv)
                source = "kwargs"
            else:
                resolved_kv_heads = max(1, self.num_heads // 4)
                source = "default"

        self.num_kv_heads = int(min(self.num_heads, max(1, resolved_kv_heads)))
        self._num_kv_heads_source = source

        # ---- sizes ----
        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))

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            model_type = str(_cfg_get(config, "model_type", "")).lower()
            is_dec = _cfg_get(
                config,
                "is_decoder",
                _cfg_get(
                    config,
                    "is_causal",
                    kwargs.get("is_causal", kwargs.get("is_decoder", False)),
                ),
            )
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            if any(model_type.startswith(p) for p in _dec_pfx):
                is_dec = True
            if src_attn is not None:
                if getattr(src_attn, "is_decoder", False) or hasattr(src_attn, "c_attn"):
                    is_dec = True
            self.is_decoder = bool(is_dec)
        self.is_nested_encoder = False

        # ---- LoRA ----
        comp_config = getattr(config, "compression_config", None) if config else None
        self.lora_rank = int(
            _cfg_get(
                comp_config,
                "lora_rank",
                _cfg_get(config, "lora_rank", kwargs.get("lora_rank", lora_rank)),
            )
        )
        self.lora_alpha = float(
            _cfg_get(
                comp_config,
                "lora_alpha",
                _cfg_get(config, "lora_alpha", kwargs.get("lora_alpha", lora_alpha)),
            )
        )
        self.freeze_base = self.lora_rank > 0

        # ---- projections ----
        self.query = LoRALinear(
            self.hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.key = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.value = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        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=self.freeze_base,
        )

        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(float(dropout_p))

        self.rotary_emb = getattr(src_attn, "rotary_emb", None)

        # Storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

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

        # Log the resolved head config
        logger.debug(
            "GQA: hidden=%d num_heads=%d num_kv_heads=%d (source=%s) kv_repeat=%d head_dim=%d",
            self.hidden_size,
            self.num_heads,
            self.num_kv_heads,
            self._num_kv_heads_source,
            self.kv_repeat,
            self.head_dim,
        )

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

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

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

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

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

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

    @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:
        from neuralops.core.adapters import UniversalAdapter

        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        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 as e:
                logger.debug("UniversalAdapter failed for %s: %s", proj, e)

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            weight, bias_data = _unpack_weight_bias(src.c_attn)
            if weight is not None:
                split_size = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split_size, :],
                    bias_data[:split_size] if bias_data is not None else None,
                )
                extracted["k"] = (
                    weight[split_size : 2 * split_size, :],
                    (bias_data[split_size : 2 * split_size] if bias_data is not None else None),
                )
                extracted["v"] = (
                    weight[2 * split_size :, :],
                    bias_data[2 * split_size :] if bias_data is not None else None,
                )

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            w_o, b_o = _unpack_weight_bias(src.c_proj)
            extracted["o"] = (w_o, b_o)

        # Nested BERT/RoBERTa
        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") and extracted["q"][0] is None:
                extracted["q"] = _unpack_weight_bias(s.query)
            if hasattr(s, "key") and extracted["k"][0] is None:
                extracted["k"] = _unpack_weight_bias(s.key)
            if hasattr(s, "value") and extracted["v"][0] is None:
                extracted["v"] = _unpack_weight_bias(s.value)

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            extracted["o"] = _unpack_weight_bias(src.output.dense)

        # Flat modules
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for n in names:
                    if hasattr(src, n):
                        w, b = _unpack_weight_bias(getattr(src, n))
                        if w is not None:
                            extracted[proj] = (w, b)
                            break

        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ) and not has_out_proj

        missing = [k for k in ("q", "k", "v") if extracted[k][0] is None]
        if not self.is_nested_encoder and extracted["o"][0] is None:
            missing.append("o")
        if missing:
            logger.warning("GQA copy_weights: missing %s", missing)

        def _orient(w, target_shape):
            if w.dim() != 2:
                return w
            if w.shape == target_shape:
                return w
            if w.shape == (target_shape[1], target_shape[0]):
                return w.t().contiguous()
            return w

        # GQA-aware pooling for K/V
        def _pool_gqa(w, target_out):
            """
            Map source K/V of shape (src_out, in) to (target_out, in).
            If src has more heads, mean-pool groups.
            If src has fewer heads, repeat each head.
            """
            src_out = w.shape[0]
            if src_out == target_out:
                return w
            if src_out > target_out and src_out % target_out == 0:
                rep = src_out // target_out
                return w.view(target_out, rep, -1).mean(dim=1)
            if src_out < target_out and target_out % src_out == 0:
                rep = target_out // src_out
                return w.repeat_interleave(rep, dim=0)
            # Fallback: principal slice
            out = torch.zeros(target_out, w.shape[1], dtype=w.dtype, device=w.device)
            m = min(target_out, src_out)
            out[:m].copy_(w[:m])
            return out

        with torch.no_grad():
            for key, target in (
                ("q", self.query),
                ("k", self.key),
                ("v", self.value),
                ("o", self.out_proj),
            ):
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        min_d = min(target.weight.shape[0], target.weight.shape[1])
                        target.weight.data.zero_()
                        target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                        if target.bias is not None:
                            target.bias.data.zero_()
                    target.freeze_base_weights()
                    continue

                if key in ("k", "v"):
                    w = _pool_gqa(w, target.weight.shape[0])
                    if b is not None:
                        b = _pool_gqa(b.unsqueeze(-1), target.weight.shape[0]).squeeze(-1)

                w = _orient(w, target.weight.shape)
                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))
                target.weight.data.zero_()
                target.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                )

                if b is not None and target.bias is not None:
                    mn = min(target.bias.numel(), b.numel())
                    target.bias.data.zero_()
                    target.bias.data[:mn].copy_(b[:mn].to(target.bias.device, target.bias.dtype))
                elif target.bias is not None:
                    target.bias.data.zero_()

                nrm = float(torch.linalg.norm(target.weight.data.float()).item())
                if nrm < 1e-4:
                    logger.warning("GQA: '%s' near-zero norm (%.2e)", key, nrm)

                target.freeze_base_weights()

        try:
            for m in (self.query, self.key, self.value, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(attention_mask, B, L, device):
        if attention_mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape == (B, L):
                return m
        elif attention_mask.dim() == 3:
            m = _to_bool(attention_mask)
            if m.shape[0] == B and m.shape[-1] == L:
                return m[:, -1, :]
        elif attention_mask.dim() == 4:
            m = _to_bool(attention_mask)
            m = m.any(dim=-2)
            if m.dim() == 3:
                m = m.squeeze(1)
            if m.shape == (B, L):
                return m
        return None

    # ------------------------------------------------------------------
    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        while isinstance(hidden_states, (tuple, list)):
            hidden_states = hidden_states[0]
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        key_value_states = kwargs.get("key_value_states")
        is_cross = key_value_states is not None or encoder_hidden_states is not None
        kv_input = (
            encoder_hidden_states
            if encoder_hidden_states is not None
            else (key_value_states if key_value_states is not None 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

        Q = self.query(hidden_states).view(B, N_q, H, d).transpose(1, 2).contiguous()
        K = self.key(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V = self.value(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # ---- RoPE ----
        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
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin, kwargs.get("position_ids"))

        # ---- GQA expand ----
        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        # ---- Attention scores ----
        scale = d**-0.5
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        # ---- Causal mask ----
        is_causal = bool(self.is_decoder and not is_cross)
        if is_causal:
            past_len = max(0, N_k - N_q)
            causal_diag = 1 + past_len
            causal = (
                torch.triu(
                    torch.full((N_q, N_k), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=causal_diag,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        # ---- Padding mask ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is None:
            eff_mask = kwargs.get("mask")
        valid_mask = self._prepare_valid_mask(eff_mask, B, N_k, device)
        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

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

        # ---- head_mask ----
        if head_mask is not None:
            hm = head_mask.to(attn.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, H, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], H, 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            attn = attn * hm

        # ---- Context ----
        context = torch.matmul(attn, V)
        context = context.transpose(1, 2).contiguous().view(B, N_q, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)

        return output, (attn if output_attentions else None)

    def get_ortho_loss(self):
        return torch.zeros((), device=self.query.weight.device)

    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("query", self.query),
            ("key", self.key),
            ("value", self.value),
            ("out_proj", self.out_proj),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["total"] = sum(out.values())
        return out

Every 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_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.key

prop k_proj

Expand source code

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

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.query

prop q_proj

Expand source code

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

prop v_lin

Expand source code

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

prop v_proj

Expand source code

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

Methods

def get_ortho_loss(self)

Expand source code

def get_ortho_loss(self):
    return torch.zeros((), device=self.query.weight.device)

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("query", self.query),
        ("key", self.key),
        ("value", self.value),
        ("out_proj", self.out_proj),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class HALOReGLAAttention (config: Any | None = None,
rank: int | None = None,
hybrid_ratio: int = 4,
pretrained_attn: nn.Module | None = None,
gate_init_bias: float = 1.8,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("halo_regla")
@ATTENTION_REGISTRY.register("halo_regla_attention")
class HALOReGLAAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        hybrid_ratio: int = 4,
        pretrained_attn: nn.Module | None = None,
        gate_init_bias: float = 1.8,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank or 64)
        self.config = config

        # Resolve dim
        dim = None
        if config is not None:
            dim = _cfg_get(
                config,
                "hidden_size",
                _cfg_get(
                    config,
                    "d_model",
                    _cfg_get(config, "n_embd", _cfg_get(config, "dim", None)),
                ),
            )
        if dim is None and pretrained_attn is not None:
            dim = getattr(
                pretrained_attn,
                "hidden_size",
                getattr(
                    pretrained_attn,
                    "embed_dim",
                    getattr(pretrained_attn, "all_head_size", None),
                ),
            )
            if dim is None:
                for cand in ("q_proj", "q_lin", "query", "c_attn"):
                    m = getattr(pretrained_attn, cand, None)
                    if m is not None and hasattr(m, "in_features"):
                        dim = m.in_features
                        break
        dim = dim or kwargs.get("hidden_size", kwargs.get("d_model", kwargs.get("dim", 768)))

        # Resolve num_heads
        num_heads = None
        if config is not None:
            num_heads = _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(
                    config,
                    "num_heads",
                    _cfg_get(config, "n_heads", _cfg_get(config, "n_head", None)),
                ),
            )
        if num_heads is None and pretrained_attn is not None:
            num_heads = getattr(
                pretrained_attn,
                "num_heads",
                getattr(pretrained_attn, "num_attention_heads", None),
            )
        num_heads = num_heads or kwargs.get("num_attention_heads", kwargs.get("num_heads", 12))

        self.dim: int = int(dim)
        self.num_heads: int = int(num_heads)
        self.num_kv_heads: int = int(
            _cfg_get(
                config,
                "num_key_value_heads",
                kwargs.get("num_key_value_heads", self.num_heads),
            )
        )
        self.head_dim: int = int(
            _cfg_get(
                config,
                "head_dim",
                kwargs.get("head_dim", self.dim // max(1, self.num_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
        self.hybrid_ratio: int = hybrid_ratio

        model_type = str(_cfg_get(config, "model_type", "")).lower()
        is_dec = _cfg_get(
            config,
            "is_decoder",
            _cfg_get(
                config,
                "is_causal",
                kwargs.get("is_causal", kwargs.get("is_decoder", False)),
            ),
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if any(_mt.startswith(p) for p in _decoder_prefixes):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)
        self.is_nested_encoder: bool = False

        comp_config = getattr(config, "compression_config", None)
        self.lora_rank: int = int(
            _cfg_get(
                comp_config,
                "lora_rank",
                _cfg_get(config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
        )
        self.lora_alpha: float = float(
            _cfg_get(
                comp_config,
                "lora_alpha",
                _cfg_get(config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        )
        self.freeze_base: bool = self.lora_rank > 0

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

        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        self.gate = nn.Linear(self.dim, self.num_heads)
        self.gate.bias.data.fill_(gate_init_bias)

        dropout_p = _cfg_get(
            config, "dropout", _cfg_get(config, "attention_probs_dropout_prob", 0.1)
        )
        self.dropout = nn.Dropout(dropout_p)

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if source_attn is not None:
            self.copy_weights(source_attn)

    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
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c_w, c_b = _unpack_weight_bias(src.c_attn)
            if c_w is not None:
                if c_w.dim() == 2 and c_w.shape[1] == 3 * self.dim:
                    c_w = c_w.t().contiguous()
                split_size = c_w.size(0) // 3
                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:
                o_w, o_b = _unpack_weight_bias(src.c_proj)
                if o_w is not None:
                    if o_w.dim() == 2 and o_w.shape[0] != self.dim and o_w.shape[1] == self.dim:
                        o_w = o_w.t().contiguous()
                    extracted["o"] = (o_w, o_b)

        # 2. 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)
                extracted[proj_key] = _unpack_weight_bias(mod)
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                extracted["o"] = _unpack_weight_bias(src.output.dense)

        # 3. Flat modules
        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)
                        cw, cb = _unpack_weight_bias(mod)
                        if cw is not None:
                            extracted[proj_key] = (cw, cb)
                            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:
            logger.debug(
                "HALOReGLAAttention partial weights extracted; missing: %s",
                missing_after,
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = torch.linalg.norm(target_mod.weight.float()).item()
                if fnorm < 1e-4:
                    logger.debug("Weight copy for %s has small norm: %.6f", proj_key, fnorm)

        if self.freeze_base:
            for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
                mod.freeze_base_weights()

    _copy_pretrained = copy_weights

    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> HALOReGLAAttention:
        return cls(config=model_config, pretrained_attn=old_attn, **kwargs)

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if encoder_hidden_states is not None:
            return hidden_states, None

        while isinstance(hidden_states, (tuple, list)):
            hidden_states = hidden_states[0]
        B, L, D = hidden_states.shape

        q = self.q_proj(hidden_states).view(B, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k = (
            self.k_proj(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )

        # Apply RoPE if provided
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # Repeat KV for GQA
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = repeat_kv(k, repeats)
            v = repeat_kv(v, repeats)

        # 1. Standard dense softmax stream
        scale = 1.0 / math.sqrt(self.head_dim)
        dense_scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        is_causal = self.is_decoder or kwargs.get("is_causal", False)
        if is_causal:
            causal_mask = torch.triu(
                torch.full((L, L), -1e4, device=dense_scores.device), diagonal=1
            )
            dense_scores = dense_scores + causal_mask.unsqueeze(0).unsqueeze(0)

        if attention_mask is not None:
            dense_scores = apply_attention_mask(dense_scores, attention_mask)

        dense_attn = F.softmax(dense_scores, dim=-1).to(hidden_states.dtype)
        dense_attn = torch.nan_to_num(dense_attn, nan=0.0)
        dense_attn = self.dropout(dense_attn)
        dense_out = torch.matmul(dense_attn, v)  # (B, H, L, d)

        # 2. Linear stream
        q_lin = F.elu(q) + 1.0
        valid_mask = get_token_valid_mask(attention_mask, B, L, hidden_states.device)
        mask_4d = valid_mask.unsqueeze(1).unsqueeze(-1).to(k.dtype)
        k_lin = (F.elu(k) + 1.0) * mask_4d
        v_masked = v * mask_4d

        if is_causal:
            kv_outer = k_lin.unsqueeze(-1) * v_masked.unsqueeze(-2)
            cum_kv = torch.cumsum(kv_outer.float(), dim=2).to(q.dtype)
            num = (q_lin.unsqueeze(-2) @ cum_kv).squeeze(-2)
            cum_k = torch.cumsum(k_lin.float(), dim=2).to(q.dtype)
            den = (q_lin * cum_k).sum(dim=-1, keepdim=True).clamp(min=1e-5)
            linear_out = (num / den) * mask_4d
        else:
            kv = torch.matmul(k_lin.transpose(-2, -1).contiguous(), v_masked)
            z = k_lin.sum(dim=2, keepdim=True)
            num = torch.matmul(q_lin, kv)
            den = torch.matmul(q_lin, z.transpose(-2, -1).contiguous()).clamp(min=1e-5)
            linear_out = (num / den) * mask_4d

        # 3. Adaptive fusion
        gate_w = (
            self.gate.weight.to(hidden_states.dtype)
            if self.gate.weight.dtype != hidden_states.dtype
            else self.gate.weight
        )
        gate_b = (
            self.gate.bias.to(hidden_states.dtype)
            if (self.gate.bias is not None and self.gate.bias.dtype != hidden_states.dtype)
            else self.gate.bias
        )
        gate_out = F.linear(hidden_states, gate_w, gate_b)
        g = (
            torch.sigmoid(gate_out).unsqueeze(1).unsqueeze(-1)
        )  # (B, 1, L, 1) or per-head (B, H, L, 1)
        if g.shape[1] != self.num_heads:
            g = torch.sigmoid(gate_out).permute(0, 2, 1).unsqueeze(-1)  # (B, H, L, 1)
        g = g.to(hidden_states.dtype)
        fused = (
            g * dense_out.to(hidden_states.dtype) + (1.0 - g) * linear_out.to(hidden_states.dtype)
        ).to(hidden_states.dtype)

        # head_mask handling
        if head_mask is not None:
            if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                fused = fused * head_mask.view(1, self.num_heads, 1, 1).to(fused.dtype)
            elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                fused = fused * head_mask.view(head_mask.size(0), self.num_heads, 1, 1).to(
                    fused.dtype
                )
            elif head_mask.dim() == 4:
                fused = fused * head_mask.to(fused.dtype)

        out = fused.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)

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

        if out_final.shape[-1] != hidden_states.shape[-1]:
            if out_final.shape[-1] < hidden_states.shape[-1]:
                out_final = F.pad(out_final, (0, hidden_states.shape[-1] - out_final.shape[-1]))
            else:
                out_final = out_final[..., : hidden_states.shape[-1]]

        return out_final, None

Every 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.

Static methods

def from_distilbert_pretrained(old_attn: nn.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> HALOReGLAAttention

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class HGRN2GatedDeltaNet (config: Any | None = None,
rank: int | None = None,
d_model: Any | None = None,
num_heads: int | None = None,
comp_config: Any | None = None,
old_attn: nn.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("hgrn2_gated_deltanet")
class HGRN2GatedDeltaNet(BaseAttention):
    r"""HGRN-2 Gated DeltaNet Hybrid Architecture.

    Combines:
      1. HGRN-2 branch: fast element-wise linear recurrence.
      2. Gated DeltaNet branch: associative matrix state with Delta-rule.
      3. Learnable fusion:
         $Y = \gamma \odot Y_{hgrn} + (1-\gamma) \odot Y_{deltanet}$
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: Any | None = None,
        num_heads: int | None = None,
        comp_config: Any | None = None,
        old_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        # ---- d_model ----
        if d_model is None:
            if isinstance(config, int):
                d_model = config
            elif config is not None:
                d_model = (
                    _cfg_get(config, "hidden_size", None)
                    or _cfg_get(config, "d_model", None)
                    or _cfg_get(config, "n_embd", None)
                    or _cfg_get(config, "dim", None)
                )
        if d_model is None and old_attn is not None:
            d_model = (
                getattr(old_attn, "hidden_size", None)
                or getattr(old_attn, "embed_dim", None)
                or getattr(old_attn, "all_head_size", None)
            )
            if d_model is None:
                for cand in ("q_proj", "q_lin", "query", "c_attn"):
                    m = getattr(old_attn, cand, None)
                    if m is not None and hasattr(m, "in_features"):
                        d_model = m.in_features
                        break
        if d_model is None:
            d_model = kwargs.get("hidden_size", kwargs.get("d_model", 768))

        # ---- num_heads ----
        if num_heads is None and config is not None and not isinstance(config, int):
            num_heads = (
                _cfg_get(config, "num_attention_heads", None)
                or _cfg_get(config, "num_heads", None)
                or _cfg_get(config, "n_heads", None)
                or _cfg_get(config, "n_head", None)
            )
        if num_heads is None and old_attn is not None:
            num_heads = (
                getattr(old_attn, "num_heads", None)
                or getattr(old_attn, "num_attention_heads", None)
                or getattr(old_attn, "n_heads", None)
            )
        if num_heads is None:
            num_heads = kwargs.get("num_attention_heads", kwargs.get("num_heads", 12))

        self.d_model = int(d_model)
        self.num_heads = int(num_heads)

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            model_type = str(_cfg_get(config, "model_type", "")).lower()
            is_dec = bool(
                _cfg_get(config, "is_decoder", False)
                or _cfg_get(config, "is_causal", False)
                or kwargs.get("is_causal", False)
                or kwargs.get("is_decoder", False)
            )
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
            )
            if any(model_type.startswith(p) for p in _dec_pfx):
                is_dec = True
            if old_attn is not None:
                if getattr(old_attn, "is_decoder", False) or hasattr(old_attn, "c_attn"):
                    is_dec = True
            self.is_decoder = bool(is_dec)

        # ---- LoRA ----
        self.lora_rank = int(_cfg_get(comp_config, "lora_rank", kwargs.get("lora_rank", lora_rank)))
        self.lora_alpha = float(
            _cfg_get(comp_config, "lora_alpha", kwargs.get("lora_alpha", lora_alpha))
        )
        self.freeze_base = self.lora_rank > 0

        # ---- branch kwargs ----
        sub_kwargs = {
            k: v
            for k, v in kwargs.items()
            if k not in ("is_decoder", "is_causal", "lora_rank", "lora_alpha")
        }

        # ---- branches ----
        self.hgrn = HGRN2Layer(
            config=config,
            rank=rank,
            d_model=self.d_model,
            is_decoder=self.is_decoder,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            **sub_kwargs,
        )
        self.deltanet = GatedDeltaNet(
            config=config,
            rank=rank,
            d_model=self.d_model,
            num_heads=self.num_heads,
            is_decoder=self.is_decoder,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            **sub_kwargs,
        )

        # ---- gate ----
        # HGD-FIX-02: balanced init
        self.gate = nn.Linear(self.d_model, self.d_model)
        with torch.no_grad():
            self.gate.weight.zero_()
            self.gate.bias.fill_(0.0)  # sigmoid(0) = 0.5

        # ---- output ----
        self.out_proj = LoRALinear(
            self.d_model,
            self.d_model,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Storage
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

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

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and old_attn is not None:
            p = next(old_attn.parameters(), None)
            if p is not None:
                target_dtype = p.dtype
        if target_dtype is not None:
            self.to(target_dtype)

        self._per_layer_trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        if self._per_layer_trainable > 5e6:
            logger.warning(
                "HGRN2GatedDeltaNet per-layer trainable: %.1fM.",
                self._per_layer_trainable / 1e6,
            )

    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        try:
            for p in old_attn.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        hgrn_ok = False
        deltanet_ok = False
        if hasattr(self.hgrn, "copy_weights"):
            try:
                self.hgrn.copy_weights(old_attn)
                hgrn_ok = True
            except Exception as ex:
                logger.warning("HGRN copy_weights failed: %s", ex)
        if hasattr(self.deltanet, "copy_weights"):
            try:
                self.deltanet.copy_weights(old_attn)
                deltanet_ok = True
            except Exception as ex:
                logger.warning("DeltaNet copy_weights failed: %s", ex)

        if not hgrn_ok and not deltanet_ok:
            logger.warning("Both HGRN and DeltaNet copy_weights failed.")
        elif hgrn_ok and not deltanet_ok:
            with torch.no_grad():
                self.gate.weight.zero_()
                self.gate.bias.fill_(math.log(0.9 / 0.1))  # ~2.20
            logger.info("Only HGRN copied — biased gate to HGRN")
        elif deltanet_ok and not hgrn_ok:
            with torch.no_grad():
                self.gate.weight.zero_()
                self.gate.bias.fill_(-math.log(0.9 / 0.1))
            logger.info("Only DeltaNet copied — biased gate to DeltaNet")

        # ---- out_proj ----
        w_o, b_o = None, None
        try:
            adapter = UniversalAdapter()
            w_o, b_o = adapter.get_weight_bias(old_attn, "o")
        except Exception:
            pass

        if w_o is None:
            for cand in (
                "out_proj",
                "out_lin",
                "o_proj",
                "output",
                "proj",
                "c_proj",
                "dense",
            ):
                for target in (
                    old_attn,
                    getattr(old_attn, "self", None),
                    getattr(old_attn, "attention", None),
                    getattr(old_attn, "output", None),
                ):
                    if target is not None and hasattr(target, cand):
                        mod = getattr(target, cand)
                        if hasattr(mod, "dense"):
                            mod = mod.dense
                        cw, cb = _unpack_weight_bias(mod)
                        if cw is not None:
                            w_o, b_o = cw, cb
                            break
                if w_o is not None:
                    break

        if w_o is not None:
            with torch.no_grad():
                tw = self.out_proj.weight
                if w_o.shape == tw.shape:
                    tw.copy_(w_o.to(tw.dtype))
                elif w_o.t().shape == tw.shape:
                    tw.copy_(w_o.t().to(tw.dtype))
                elif w_o.numel() == tw.numel():
                    tw.copy_(w_o.reshape(tw.shape).to(tw.dtype))
                else:
                    min_r = min(tw.shape[0], w_o.shape[0])
                    min_c = min(tw.shape[1], w_o.shape[1])
                    tw.data.zero_()
                    tw.data[:min_r, :min_c].copy_(w_o[:min_r, :min_c].to(tw.dtype))
                    logger.warning(
                        "Out proj shape mismatch: %s → %s.",
                        tuple(w_o.shape),
                        tuple(tw.shape),
                    )

                if b_o is not None and self.out_proj.bias is not None:
                    tb = self.out_proj.bias
                    if b_o.shape == tb.shape:
                        tb.copy_(b_o.to(tb.dtype))
                    elif b_o.numel() == tb.numel():
                        tb.copy_(b_o.reshape(tb.shape).to(tb.dtype))
                    else:
                        min_b = min(tb.numel(), b_o.numel())
                        tb.data.zero_()
                        tb.data[:min_b].copy_(b_o[:min_b].to(tb.dtype))
        else:
            logger.warning(
                "Out projection weight not found for %s.",
                type(old_attn).__name__,
            )

        if self.freeze_base:
            self.out_proj.freeze_base_weights()
            self.gate.weight.requires_grad = False
            if self.gate.bias is not None:
                self.gate.bias.requires_grad = False

        # Storage accounting
        try:
            for m in (self.hgrn, self.deltanet, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

    # ------------------------------------------------------------------
    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]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("HGRN2GatedDeltaNet does not support cross-attention.")

        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        hgrn_res = self.hgrn(
            hidden_states,
            attention_mask=attention_mask,
            head_mask=head_mask,
            **kwargs,
        )
        hgrn_out = hgrn_res[0] if isinstance(hgrn_res, tuple) else hgrn_res

        delta_res = self.deltanet(
            hidden_states,
            attention_mask=attention_mask,
            head_mask=head_mask,
            **kwargs,
        )
        delta_out = delta_res[0] if isinstance(delta_res, tuple) else delta_res

        if hgrn_out.shape != delta_out.shape:
            raise RuntimeError(
                f"HGRN2GatedDeltaNet branch outputs mismatch: "
                f"{tuple(hgrn_out.shape)} vs {tuple(delta_out.shape)}."
            )

        gate_dtype = next(self.gate.parameters()).dtype
        gamma = torch.sigmoid(self.gate(hidden_states.to(gate_dtype))).to(hidden_states.dtype)
        fused = gamma * hgrn_out.to(hidden_states.dtype) + (1.0 - gamma) * delta_out.to(
            hidden_states.dtype
        )
        out = self.out_proj(fused)

        if out.shape[-1] != hidden_states.shape[-1]:
            raise RuntimeError(
                f"HGRN2GatedDeltaNet output last dim {out.shape[-1]} "
                f"!= hidden {hidden_states.shape[-1]}."
            )

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

    def get_ortho_loss(self):
        return torch.zeros((), device=self.out_proj.weight.device)

    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("hgrn", self.hgrn),
            ("deltanet", self.deltanet),
            ("out_proj", self.out_proj),
            ("gate", self.gate),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["total"] = sum(out.values())
        return out

HGRN-2 Gated DeltaNet Hybrid Architecture.

Combines -----= 1. HGRN-2 branch: fast element-wise linear recurrence. 2. Gated DeltaNet branch: associative matrix state with Delta-rule. 3. Learnable fusion: $Y = \gamma \odot Y_{hgrn} + (1-\gamma) \odot Y_{deltanet}$

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.

Methods

def get_ortho_loss(self)

Expand source code

def get_ortho_loss(self):
    return torch.zeros((), device=self.out_proj.weight.device)

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("hgrn", self.hgrn),
        ("deltanet", self.deltanet),
        ("out_proj", self.out_proj),
        ("gate", self.gate),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class HGRN2Layer (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
res_clip: float = 30.0,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
freeze_base: bool = True,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("hgrn2_layer")
@ATTENTION_REGISTRY.register("hgrn2")
class HGRN2Layer(BaseAttention):
    r"""Hierarchically Gated Recurrent Network v2 (HGRN-2) Layer.
    Supports causal mode for decoders, bidirectional mode for encoders, LoRA adapters,
    and pretrained weight copying from DistilBERT, RoBERTa, GPT-2, and Qwen3.5.

    NOTE: This is NOT a compression module. It replaces attention with an RNN.
    Parameter count is the same. Only use for speed improvements on
    long sequences (when parallel scan is implemented).
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        res_clip: float = 30.0,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        freeze_base: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        # Resolve d_model
        if d_model is None:
            d_model = kwargs.get("d_model") or kwargs.get("hidden_size")
            if d_model is None:
                if isinstance(config, int):
                    d_model = config
                elif config is not None:
                    d_model = _cfg_get(
                        config,
                        "hidden_size",
                        _cfg_get(
                            config,
                            "d_model",
                            _cfg_get(config, "n_embd", _cfg_get(config, "dim", None)),
                        ),
                    )
            if d_model is None:
                d_model = kwargs.get("dim", 128)

        if hasattr(d_model, "hidden_size"):
            d_model = d_model.hidden_size

        self.d_model: int = int(d_model)
        self.dim: int = self.d_model
        self.hidden_size: int = self.d_model
        self.res_clip: float = res_clip
        self.freeze_base: bool = freeze_base

        _mt = str(_cfg_get(config, "model_type", "") or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = bool(
                _cfg_get(config, "is_decoder", False)
                or _cfg_get(config, "causal", False)
                or any(_mt.startswith(p) for p in _decoder_prefixes)
                or kwargs.get("is_decoder", False)
                or kwargs.get("is_causal", False)
            )

        self.is_nested_encoder: bool = False
        self.lora_rank = lora_rank if lora_rank is not None else 16
        self.lora_alpha = lora_alpha
        self.lora_scaling = (
            (self.lora_alpha / max(1, self.lora_rank)) if self.lora_rank > 0 else 1.0
        )

        self.forget_gate = nn.Linear(self.d_model, self.d_model, bias=True)
        self.input_gate = nn.Linear(self.d_model, self.d_model, bias=True)
        self.v_proj = nn.Linear(self.d_model, self.d_model, bias=False)
        self.g_proj = nn.Linear(self.d_model, self.d_model, bias=True)
        self.state_norm = nn.LayerNorm(self.d_model)
        self.out_proj = nn.Linear(self.d_model, self.d_model, bias=True)

        # Initialize gates with sensible defaults for long-term retention and unity gain
        nn.init.normal_(self.forget_gate.weight, mean=0.0, std=0.01)
        nn.init.constant_(self.forget_gate.bias, -2.5)
        nn.init.normal_(self.input_gate.weight, mean=0.0, std=0.02)
        nn.init.constant_(self.input_gate.bias, 0.0)
        nn.init.normal_(self.g_proj.weight, mean=0.0, std=0.02)
        nn.init.constant_(self.g_proj.bias, 1.28)
        nn.init.normal_(self.v_proj.weight, mean=0.0, std=0.02)

        # LoRA Adapters
        if self.lora_rank > 0:
            self.v_lora_A = nn.Parameter(torch.randn(self.lora_rank, self.d_model) * 0.01)
            self.v_lora_B = nn.Parameter(torch.zeros(self.d_model, self.lora_rank))
            self.out_lora_A = nn.Parameter(torch.randn(self.lora_rank, self.d_model) * 0.01)
            self.out_lora_B = nn.Parameter(torch.zeros(self.d_model, self.lora_rank))
        else:
            self.register_parameter("v_lora_A", None)
            self.register_parameter("v_lora_B", None)
            self.register_parameter("out_lora_A", None)
            self.register_parameter("out_lora_B", None)

        if self.freeze_base:
            for mod in (self.v_proj, self.out_proj):
                for p in mod.parameters():
                    p.requires_grad = False
            for mod in (self.forget_gate, self.input_gate, self.g_proj):
                for p in mod.parameters():
                    p.requires_grad = True

        old_attn = kwargs.get("old_attn", kwargs.get("pretrained_attn"))
        if old_attn is not None:
            self.copy_weights(old_attn)
            target_dtype = kwargs.get("dtype")
            if target_dtype is None:
                p = next(old_attn.parameters(), None)
                if p is not None and isinstance(p.dtype, torch.dtype):
                    target_dtype = p.dtype
            if target_dtype is not None:
                self.to(target_dtype)

    def copy_weights(self, old_attn: Any) -> None:
        """Transfer pretrained projection weights across all 4 architectures."""
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        is_nested = (
            src_cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            )
            or (hasattr(old_attn, "self") and hasattr(old_attn, "output"))
        )
        self.is_nested_encoder = bool(is_nested)

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o, b_o = _unpack_weight_bias(old_attn.output.dense)
        else:
            src = old_attn

        # Pattern 0: PyTorch MultiheadAttention
        if hasattr(src, "in_proj_weight") and src.in_proj_weight is not None:
            in_w, in_b = _unpack_weight_bias(src)
            if in_w is None:
                in_w = getattr(src, "in_proj_weight", None)
                in_b = getattr(src, "in_proj_bias", None)
            if in_w is not None:
                chunks = in_w.chunk(3, dim=0)
                w_q, w_k, w_v = chunks[0], chunks[1], chunks[2]
            if in_b is not None:
                b_chunks = in_b.chunk(3, dim=0)
                b_q, b_k, b_v = b_chunks[0], b_chunks[1], b_chunks[2]
            if hasattr(src, "out_proj"):
                w_o, b_o = _unpack_weight_bias(src.out_proj)

        # Pattern 1: GPT-2 Fused c_attn
        elif hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
            if c_attn_w is not None:
                split_size = c_attn_w.shape[0] // 3
                w_q = c_attn_w[:split_size, :]
                w_k = c_attn_w[split_size : 2 * split_size, :]
                w_v = c_attn_w[2 * split_size :, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:split_size]
                    b_k = c_attn_b[split_size : 2 * split_size]
                    b_v = c_attn_b[2 * split_size :]
            if hasattr(src, "c_proj"):
                w_o, b_o = _unpack_weight_bias(src.c_proj)

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q, b_q = _unpack_weight_bias(src.query)
            w_k, b_k = _unpack_weight_bias(src.key)
            w_v, b_v = _unpack_weight_bias(src.value)
            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q, b_q = _unpack_weight_bias(src.q_proj)
            w_k, b_k = _unpack_weight_bias(src.k_proj)
            w_v, b_v = _unpack_weight_bias(src.v_proj)
            if hasattr(src, "o_proj"):
                w_o, b_o = _unpack_weight_bias(src.o_proj)

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q, b_q = _unpack_weight_bias(src.q_lin)
            w_k, b_k = _unpack_weight_bias(src.k_lin)
            w_v, b_v = _unpack_weight_bias(src.v_lin)
            if hasattr(src, "out_lin"):
                w_o, b_o = _unpack_weight_bias(src.out_lin)

        # Fallback: UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter(old_attn)
                extracted = adapter.extract_qkv_out()
                if extracted:
                    w_q = extracted.get("q_weight")
                    extracted.get("q_bias")
                    w_k = extracted.get("k_weight")
                    extracted.get("k_bias")
                    w_v = extracted.get("v_weight")
                    b_v = extracted.get("v_bias")
                    if w_o is None:
                        w_o = extracted.get("out_weight")
                        b_o = extracted.get("out_bias")
            except Exception as ex:
                logger.debug("UniversalAdapter fallback failed: %s", ex)

        def _copy_param(target_linear: nn.Linear, w: torch.Tensor | None, b: torch.Tensor | None):
            if w is None:
                return
            with torch.no_grad():
                tw = target_linear.weight
                if w.shape == tw.shape:
                    tw.copy_(w.to(tw.dtype))
                else:
                    min_r = min(tw.shape[0], w.shape[0])
                    min_c = min(tw.shape[1], w.shape[1])
                    tw.data.zero_()
                    tw.data[:min_r, :min_c].copy_(w[:min_r, :min_c].to(tw.dtype))
                    logger.warning(
                        "Partial weight copy: source %s → target %s",
                        tuple(w.shape),
                        tuple(tw.shape),
                    )
                if b is not None and target_linear.bias is not None:
                    tb = target_linear.bias
                    if b.shape == tb.shape:
                        tb.copy_(b.to(tb.dtype))
                    else:
                        min_b = min(tb.numel(), b.numel())
                        tb.data.zero_()
                        tb.data[:min_b].copy_(b[:min_b].to(tb.dtype))

        # Copy V -> v_proj and O -> out_proj across all architectures
        _copy_param(self.v_proj, w_v, b_v)
        _copy_param(self.out_proj, w_o, b_o)

        fnorm = float(torch.linalg.norm(self.v_proj.weight.detach().float()))
        if fnorm < 1e-4:
            logger.warning("HGRN2Layer v_proj weight norm is near zero: %g", fnorm)

        if self.freeze_base:
            for mod in (self.v_proj, self.out_proj):
                for p in mod.parameters():
                    p.requires_grad = False
            for mod in (self.forget_gate, self.input_gate, self.g_proj):
                for p in mod.parameters():
                    p.requires_grad = True

    @property
    def query(self) -> nn.Module:
        return self.g_proj

    @property
    def key(self) -> nn.Module:
        return self.input_gate

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

    def _scan_direction(
        self,
        f: torch.Tensor,
        iv: torch.Tensor,
        chunk_size: int = 32,
    ) -> torch.Tensor:
        # --- Triton fast path ---
        if HAS_TRITON_HGRN2 and hgrn2_scan_triton is not None and f.is_cuda:
            try:
                out = hgrn2_scan_triton(f.contiguous(), iv.contiguous(), res_clip=self.res_clip)
                logger.debug("HGRN2: Triton scan used")
                return out
            except Exception as e:
                logger.warning("HGRN2 Triton scan failed, falling back: %s", e)
                logger.debug("HGRN2: Triton FAILED: %s", e)

        # --- PyTorch fallback ---
        logger.debug("HGRN2: PyTorch fallback")
        B, L, D = iv.shape
        f_fp32 = f.float()
        iv_fp32 = iv.float()
        log_f = torch.log(f_fp32.clamp(min=1e-6))

        h_chunks = []
        h_prev = torch.zeros(B, D, device=iv.device, dtype=torch.float32)

        for start in range(0, L, chunk_size):
            end = min(start + chunk_size, L)
            C = end - start
            f_c = log_f[:, start:end]
            iv_c = iv_fp32[:, start:end]
            A = torch.cumsum(f_c, dim=1)
            A_diff = A.unsqueeze(2) - A.unsqueeze(1)
            mask = torch.tril(torch.ones(C, C, device=iv.device, dtype=torch.bool))
            A_diff = torch.where(
                mask.unsqueeze(0).unsqueeze(-1), A_diff, torch.full_like(A_diff, -1e4)
            )
            decay_weights = torch.exp(A_diff) * mask.unsqueeze(0).unsqueeze(-1)
            intra = torch.sum(decay_weights * iv_c.unsqueeze(1), dim=2)
            inter = h_prev.unsqueeze(1) * torch.exp(A)
            h_c = inter + intra
            if self.res_clip is not None and self.res_clip > 0:
                h_c = torch.clamp(h_c, -self.res_clip, self.res_clip)
            h_chunks.append(h_c)
            h_prev = h_c[:, -1]

        return torch.cat(h_chunks, dim=1).to(iv.dtype)

    def _recurrent_scan_parallel(
        self,
        f: torch.Tensor,
        iv: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        # mask preprocessing
        if mask is not None:
            if mask.dim() == 2:
                mask_float = mask.float()
            elif mask.dim() == 4:
                v = (mask > -1e4) if mask.dtype.is_floating_point else mask.bool()
                v = v.any(dim=1)
                if v.dim() == 3:
                    mask_float = v.any(dim=1).float()
                else:
                    mask_float = v.float()
            else:
                mask_float = mask.float()
            iv = iv * mask_float.unsqueeze(-1)

        # --- Triton path (no checkpoint, kernel handles backward) ---
        if HAS_TRITON_HGRN2 and f.is_cuda:
            h_fwd = self._scan_direction(f, iv)
            if self.is_decoder:
                return h_fwd, None
            f_rev = torch.flip(f, dims=[1])
            iv_rev = torch.flip(iv, dims=[1])
            h_rev = self._scan_direction(f_rev, iv_rev)
            h_bwd = torch.flip(h_rev, dims=[1])
            return h_fwd, h_bwd

        # --- PyTorch path with checkpointing ---
        if self.training:
            import torch.utils.checkpoint as cp

            h_fwd = cp.checkpoint(self._scan_direction, f, iv, use_reentrant=False)
        else:
            h_fwd = self._scan_direction(f, iv)
        if self.is_decoder:
            return h_fwd, None
        f_rev = torch.flip(f, dims=[1])
        iv_rev = torch.flip(iv, dims=[1])
        if self.training:
            import torch.utils.checkpoint as cp

            h_rev = cp.checkpoint(self._scan_direction, f_rev, iv_rev, use_reentrant=False)
        else:
            h_rev = self._scan_direction(f_rev, iv_rev)
        h_bwd = torch.flip(h_rev, dims=[1])
        return h_fwd, h_bwd

    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]:
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        if encoder_hidden_states is not None:
            raise NotImplementedError("Cross-attention is not supported in HGRN2Layer")

        B, L, D = hidden_states.shape

        target_dtype = hidden_states.dtype

        # 1. Compute gates and values with bounded log decay to prevent gate saturation
        raw_log_decay = -F.softplus(self.forget_gate(hidden_states))
        bounded_log_decay = raw_log_decay.clamp(min=-3.0, max=-1e-4)
        f = torch.exp(bounded_log_decay).to(target_dtype)

        # Input modulation: tie input scale with (1 - f) to prevent recurrence state sign-flipping
        i = torch.sigmoid(self.input_gate(hidden_states)).to(target_dtype)
        input_scale = ((1.0 - f) * i).to(target_dtype)
        v = self.v_proj(hidden_states).to(target_dtype)
        if self.lora_rank > 0 and self.v_lora_A is not None and self.v_lora_B is not None:
            v_lora = (
                torch.matmul(
                    torch.matmul(hidden_states, self.v_lora_A.t().contiguous().to(target_dtype)),
                    self.v_lora_B.t().contiguous().to(target_dtype),
                )
                * self.lora_scaling
            )
            v = v + v_lora

        z = self.g_proj(hidden_states).to(target_dtype)
        iv = (input_scale * v).to(target_dtype)

        # 2. Mask processing
        mask_bool = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                mask_bool = attention_mask.bool()
            elif attention_mask.dim() == 4:
                mask_bool = (attention_mask[:, 0, 0, :] > -10000.0).bool()

        # 3. Recurrent Scan
        h_fwd, h_bwd = self._recurrent_scan_parallel(f, iv, mask_bool)

        if self.is_decoder or h_bwd is None:
            h_ctx = h_fwd.to(target_dtype)
        else:
            h_ctx = (h_fwd + h_bwd).to(target_dtype)

        # Context stabilization with token self-value preservation
        sn_dtype = (
            self.state_norm.weight.dtype
            if hasattr(self.state_norm, "weight") and self.state_norm.weight is not None
            else target_dtype
        )
        h_norm = self.state_norm(h_ctx.to(sn_dtype)).to(target_dtype)
        gated = ((0.5 * v + 0.5 * h_norm) * F.silu(z)).to(target_dtype)
        out = self.out_proj(gated).to(target_dtype)
        if self.lora_rank > 0 and self.out_lora_A is not None and self.out_lora_B is not None:
            out_lora = (
                torch.matmul(
                    torch.matmul(gated, self.out_lora_A.t().contiguous().to(target_dtype)),
                    self.out_lora_B.t().contiguous().to(target_dtype),
                )
                * self.lora_scaling
            )
            out = out + out_lora

        if out.shape[-1] != hidden_states.shape[-1]:
            if out.shape[-1] < hidden_states.shape[-1]:
                out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
            else:
                out = out[..., : hidden_states.shape[-1]]

        if attention_mask is not None and mask_bool is not None:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        return out.to(target_dtype), None

Hierarchically Gated Recurrent Network v2 (HGRN-2) Layer. Supports causal mode for decoders, bidirectional mode for encoders, LoRA adapters, and pretrained weight copying from DistilBERT, RoBERTa, GPT-2, and Qwen3.5.

NOTE: This is NOT a compression module. It replaces attention with an RNN. Parameter count is the same. Only use for speed improvements on long sequences (when parallel scan is implemented).

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.

Instance variables

prop key : nn.Module

Expand source code

@property
def key(self) -> nn.Module:
    return self.input_gate

prop query : nn.Module

Expand source code

@property
def query(self) -> nn.Module:
    return self.g_proj

prop value : nn.Module

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Transfer pretrained projection weights across all 4 architectures."""
    w_q, b_q = None, None
    w_k, b_k = None, None
    w_v, b_v = None, None
    w_o, b_o = None, None

    src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
    is_nested = (
        src_cls_name in ("BertSelfAttention", "RobertaSelfAttention")
        or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        )
        or (hasattr(old_attn, "self") and hasattr(old_attn, "output"))
    )
    self.is_nested_encoder = bool(is_nested)

    if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
        src = old_attn.self
        if hasattr(old_attn.output, "dense"):
            w_o, b_o = _unpack_weight_bias(old_attn.output.dense)
    else:
        src = old_attn

    # Pattern 0: PyTorch MultiheadAttention
    if hasattr(src, "in_proj_weight") and src.in_proj_weight is not None:
        in_w, in_b = _unpack_weight_bias(src)
        if in_w is None:
            in_w = getattr(src, "in_proj_weight", None)
            in_b = getattr(src, "in_proj_bias", None)
        if in_w is not None:
            chunks = in_w.chunk(3, dim=0)
            w_q, w_k, w_v = chunks[0], chunks[1], chunks[2]
        if in_b is not None:
            b_chunks = in_b.chunk(3, dim=0)
            b_q, b_k, b_v = b_chunks[0], b_chunks[1], b_chunks[2]
        if hasattr(src, "out_proj"):
            w_o, b_o = _unpack_weight_bias(src.out_proj)

    # Pattern 1: GPT-2 Fused c_attn
    elif hasattr(src, "c_attn") and hasattr(src, "c_proj"):
        c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
        if c_attn_w is not None:
            split_size = c_attn_w.shape[0] // 3
            w_q = c_attn_w[:split_size, :]
            w_k = c_attn_w[split_size : 2 * split_size, :]
            w_v = c_attn_w[2 * split_size :, :]
            if c_attn_b is not None:
                b_q = c_attn_b[:split_size]
                b_k = c_attn_b[split_size : 2 * split_size]
                b_v = c_attn_b[2 * split_size :]
        if hasattr(src, "c_proj"):
            w_o, b_o = _unpack_weight_bias(src.c_proj)

    # Pattern 2: RoBERTa / BERT separate query, key, value
    elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
        w_q, b_q = _unpack_weight_bias(src.query)
        w_k, b_k = _unpack_weight_bias(src.key)
        w_v, b_v = _unpack_weight_bias(src.value)
        if w_o is None:
            for attr in ("dense", "out_proj", "o_proj"):
                if hasattr(src, attr):
                    w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                    break

    # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
    elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
        w_q, b_q = _unpack_weight_bias(src.q_proj)
        w_k, b_k = _unpack_weight_bias(src.k_proj)
        w_v, b_v = _unpack_weight_bias(src.v_proj)
        if hasattr(src, "o_proj"):
            w_o, b_o = _unpack_weight_bias(src.o_proj)

    # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
    elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
        w_q, b_q = _unpack_weight_bias(src.q_lin)
        w_k, b_k = _unpack_weight_bias(src.k_lin)
        w_v, b_v = _unpack_weight_bias(src.v_lin)
        if hasattr(src, "out_lin"):
            w_o, b_o = _unpack_weight_bias(src.out_lin)

    # Fallback: UniversalAdapter
    if w_q is None or w_k is None or w_v is None:
        try:
            adapter = UniversalAdapter(old_attn)
            extracted = adapter.extract_qkv_out()
            if extracted:
                w_q = extracted.get("q_weight")
                extracted.get("q_bias")
                w_k = extracted.get("k_weight")
                extracted.get("k_bias")
                w_v = extracted.get("v_weight")
                b_v = extracted.get("v_bias")
                if w_o is None:
                    w_o = extracted.get("out_weight")
                    b_o = extracted.get("out_bias")
        except Exception as ex:
            logger.debug("UniversalAdapter fallback failed: %s", ex)

    def _copy_param(target_linear: nn.Linear, w: torch.Tensor | None, b: torch.Tensor | None):
        if w is None:
            return
        with torch.no_grad():
            tw = target_linear.weight
            if w.shape == tw.shape:
                tw.copy_(w.to(tw.dtype))
            else:
                min_r = min(tw.shape[0], w.shape[0])
                min_c = min(tw.shape[1], w.shape[1])
                tw.data.zero_()
                tw.data[:min_r, :min_c].copy_(w[:min_r, :min_c].to(tw.dtype))
                logger.warning(
                    "Partial weight copy: source %s → target %s",
                    tuple(w.shape),
                    tuple(tw.shape),
                )
            if b is not None and target_linear.bias is not None:
                tb = target_linear.bias
                if b.shape == tb.shape:
                    tb.copy_(b.to(tb.dtype))
                else:
                    min_b = min(tb.numel(), b.numel())
                    tb.data.zero_()
                    tb.data[:min_b].copy_(b[:min_b].to(tb.dtype))

    # Copy V -> v_proj and O -> out_proj across all architectures
    _copy_param(self.v_proj, w_v, b_v)
    _copy_param(self.out_proj, w_o, b_o)

    fnorm = float(torch.linalg.norm(self.v_proj.weight.detach().float()))
    if fnorm < 1e-4:
        logger.warning("HGRN2Layer v_proj weight norm is near zero: %g", fnorm)

    if self.freeze_base:
        for mod in (self.v_proj, self.out_proj):
            for p in mod.parameters():
                p.requires_grad = False
        for mod in (self.forget_gate, self.input_gate, self.g_proj):
            for p in mod.parameters():
                p.requires_grad = True

Transfer pretrained projection weights across all 4 architectures.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class HGRN2OuterBlock (config: Any | None = None,
rank: int | None = None,
dim: int | Any | None = None,
num_heads: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
chunk_size: int = 64,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "hgrn2_outer", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
class HGRN2OuterBlock(BaseAttention):
    supports_causal: bool = True
    supports_bidirectional: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        dim: int | Any | None = None,
        num_heads: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        chunk_size: int = 64,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        # ---- dim resolution ----
        if dim is None:
            if isinstance(config, int):
                dim = config
            elif config is not None:
                dim = getattr(
                    config,
                    "hidden_size",
                    getattr(
                        config,
                        "d_model",
                        getattr(config, "n_embd", getattr(config, "dim", None)),
                    ),
                )
        if dim is None:
            dim = kwargs.get("hidden_size", kwargs.get("d_model", 768))
        if hasattr(dim, "hidden_size"):
            dim = dim.hidden_size

        # ---- num_heads ----
        if num_heads is None:
            if config is not None and not isinstance(config, int):
                num_heads = getattr(
                    config,
                    "num_attention_heads",
                    getattr(
                        config,
                        "num_heads",
                        getattr(config, "n_heads", getattr(config, "n_head", None)),
                    ),
                )
        if num_heads is None:
            num_heads = kwargs.get("num_attention_heads", kwargs.get("num_heads", 12))

        self.dim = int(dim)
        self.num_heads = int(num_heads)
        while self.num_heads > 1 and (self.dim % self.num_heads != 0):
            self.num_heads -= 1
        self.head_dim = self.dim // max(1, self.num_heads)
        self.lora_rank = int(lora_rank)
        self.lora_alpha = float(lora_alpha)
        self.chunk_size = int(chunk_size)

        cfg_kv = None
        if config is not None and not isinstance(config, int):
            cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        self.num_kv_heads = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.num_key_value_groups = max(1, self.num_heads // max(1, self.num_kv_heads))

        self.q_head_size = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(getattr(config, "model_type", "") or "").lower() if config else ""
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            self.is_decoder = (
                bool(getattr(config, "is_decoder", False) if config else False)
                or bool(getattr(config, "causal", False) if config else False)
                or any(_mt.startswith(p) for p in _dec_pfx)
                or bool(kwargs.get("is_causal", False))
                or bool(kwargs.get("is_decoder", False))
            )

        self.is_nested_encoder: bool = False

        # ---- projections ----
        self.q_proj = LoRALinear(
            self.dim,
            self.q_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.W_f = nn.Linear(self.dim, self.num_heads, bias=True)
        self.state_norm = nn.LayerNorm(self.head_dim)
        self.out_proj = LoRALinear(
            self.q_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        with torch.no_grad():
            nn.init.normal_(self.W_f.weight, mean=0.0, std=0.01)
            nn.init.constant_(self.W_f.bias, -2.5)

        # ---- storage ----
        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 0

        old_attn = kwargs.get("old_attn", kwargs.get("pretrained_attn"))
        if old_attn is not None:
            self.copy_weights(old_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 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 c_proj(self):
        return self.out_proj

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

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            cls_name = old_attn.__class__.__name__
            if cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # GPT-2 fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                cls_c = src.c_attn.__class__.__name__
                if (
                    hasattr(src.c_attn, "nf")
                    or "Conv1D" in cls_c
                    or (c_attn_w.shape[0] == self.dim and c_attn_w.shape[1] != self.dim)
                ):
                    c_attn_w = c_attn_w.t()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size
                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

                out_w = src.c_proj.weight.detach()
                cls_p = src.c_proj.__class__.__name__
                if (
                    hasattr(src.c_proj, "nf")
                    or "Conv1D" in cls_p
                    or (out_w.ndim == 2 and out_w.shape[0] != self.dim)
                ):
                    out_w = out_w.t()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # BERT / RoBERTa nested
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Qwen / LLaMA
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # DistilBERT
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # UniversalAdapter fallback
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    w_q, b_q = adapter.get_weight_bias(src, "q")
                    w_k, b_k = adapter.get_weight_bias(src, "k")
                    w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            missing = []
            if w_q is None:
                missing.append("Q")
            if w_k is None:
                missing.append("K")
            if w_v is None:
                missing.append("V")
            if w_o is None and not self.is_nested_encoder:
                missing.append("O")
            if missing:
                logger.warning(
                    "HGRN2OuterBlock: missing weights %s from %s",
                    missing,
                    type(old_attn).__name__,
                )
                return

            # Storage tracking
            try:
                for p in old_attn.parameters():
                    self._orig_storage_bytes += p.numel() * p.element_size()
            except Exception:
                pass

            def _copy_to(target_mod: nn.Module, w_src: torch.Tensor, b_src):
                tw = target_mod.weight
                w_cast = w_src.to(tw.dtype)
                if w_cast.shape == tw.shape:
                    tw.data.copy_(w_cast)
                else:
                    min_r = min(tw.shape[0], w_cast.shape[0])
                    min_c = min(tw.shape[1], w_cast.shape[1])
                    tw.data.zero_()
                    tw.data[:min_r, :min_c].copy_(w_cast[:min_r, :min_c])
                if b_src is not None and target_mod.bias is not None:
                    tb = target_mod.bias
                    b_cast = b_src.to(tb.dtype)
                    mb = min(tb.numel(), b_cast.numel())
                    tb.data.zero_()
                    tb.data[:mb].copy_(b_cast[:mb])

            _copy_to(self.q_proj, w_q, b_q)
            _copy_to(self.k_proj, w_k, b_k)
            _copy_to(self.v_proj, w_v, b_v)

            if w_o is not None:
                _copy_to(self.out_proj, w_o, b_o)

            # Norm checks
            for name, mod in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                n = float(torch.linalg.norm(mod.weight.float()))
                if n < 1e-4:
                    logger.warning("HGRN2OuterBlock: %s weight near-zero norm (%.2e)", name, n)

            try:
                for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
                    for p in m.parameters():
                        self._comp_storage_bytes += p.numel() * p.element_size()
            except Exception:
                pass

    # ------------------------------------------------------------------
    # HGRN2-FIX-01: chunked scan (vectorized)
    # ------------------------------------------------------------------
    def _chunked_scan(
        self,
        q: torch.Tensor,  # (B, H, L, d)
        k: torch.Tensor,  # (B, H, L, d)
        v: torch.Tensor,  # (B, H, L, d)
        f: torch.Tensor,  # (B, H, L)
        chunk_size: int = 64,
    ) -> torch.Tensor:
        """
        Chunked causal scan for outer-product recurrent linear attention.

        Recurrence:
            S_t = f_t * S_{t-1} + (1 - f_t) * v_t k_t^T
            z_t = f_t * z_{t-1} + (1 - f_t) * k_t
            y_t = (q_t @ S_t) / (q_t @ z_t)

        Within a chunk of size C, both numerator and denominator are computed
        in closed form via a causal decay matrix:
            decay[i, j] = exp(log_F[j] - log_F_prev[i])
        where log_F[j] = sum_{s<=j} log f_s within the chunk.
        """
        B, H, L, d = q.shape
        num_chunks = (L + chunk_size - 1) // chunk_size

        S_prev = torch.zeros(B, H, d, d, device=q.device, dtype=q.dtype)
        z_prev = torch.zeros(B, H, d, device=q.device, dtype=q.dtype)

        outputs = []

        for c in range(num_chunks):
            start = c * chunk_size
            end = min(start + chunk_size, L)
            C = end - start

            q_c = q[:, :, start:end]  # (B, H, C, d)
            k_c = k[:, :, start:end]
            v_c = v[:, :, start:end]
            f_c = f[:, :, start:end]  # (B, H, C)

            # ---- within-chunk cumulative log-deccay ----
            log_f = torch.log(f_c.clamp_min(1e-6))  # (B, H, C)
            log_F_c = torch.cumsum(log_f, dim=-1)  # (B, H, C)

            # Extended prev: [0, log_F_c[0], ..., log_F_c[C-1]]
            log_F_prev = torch.cat(
                [
                    torch.zeros(B, H, 1, device=q.device, dtype=log_F_c.dtype),
                    log_F_c,
                ],
                dim=-1,
            )  # (B, H, C+1)

            # log_decay[i, j] = log_F_c[j] - log_F_prev[i]
            # i = 0 → previous-chunk state
            # i >= 1 → position i-1 in chunk
            log_decay = log_F_c.unsqueeze(-2) - log_F_prev.unsqueeze(-1)  # (B, H, C+1, C)

            # Mask: i=0 always valid; i>=1 valid iff i-1 <= j  (i <= j+1)
            i_idx = torch.arange(C + 1, device=q.device).view(C + 1, 1)
            j_idx = torch.arange(C, device=q.device).view(1, C)
            valid = (i_idx == 0) | (i_idx <= j_idx + 1)
            fill_val = -1e4 if log_decay.dtype == torch.float16 else -1e9
            log_decay = log_decay.masked_fill(~valid, fill_val)
            decay = torch.exp(log_decay).to(q.dtype)  # (B, H, C+1, C)

            # ---- within-chunk write contributions ----
            one_minus_f = (1.0 - f_c).to(q.dtype)  # (B, H, C)
            W_kv = one_minus_f[..., None, None] * (
                k_c.unsqueeze(-1) * v_c.unsqueeze(-2)
            )  # (B, H, C, d, d)
            W_k = one_minus_f[..., None] * k_c  # (B, H, C, d)

            # Combined with prev state as index 0
            combined_kv = torch.cat([S_prev.unsqueeze(2), W_kv], dim=2).to(
                q.dtype
            )  # (B, H, C+1, d, d)
            combined_k = torch.cat([z_prev.unsqueeze(2), W_k], dim=2).to(q.dtype)  # (B, H, C+1, d)

            # S_all[j] = sum_i decay[i, j] * combined[i]
            S_all = torch.einsum("bhij,bhiae->bhjae", decay, combined_kv)  # (B, H, C, d, d)
            z_all = torch.einsum("bhij,bhia->bhja", decay, combined_k)  # (B, H, C, d)

            # y_j = (q_j @ S_j) / (q_j @ z_j)
            num = torch.einsum("bhjd,bhjde->bhje", q_c, S_all)  # (B, H, C, d)
            denom = torch.einsum("bhjd,bhjd->bhj", q_c, z_all).clamp_min(1e-6)  # (B, H, C)
            y_c = num / denom.unsqueeze(-1)

            outputs.append(y_c)
            S_prev = S_all[:, :, -1]  # (B, H, d, d)
            z_prev = z_all[:, :, -1]  # (B, H, d)

        return torch.cat(outputs, dim=2).to(q.dtype)  # (B, H, L, d)

    # ------------------------------------------------------------------
    # Storage
    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        return {
            "orig": self._orig_storage_bytes,
            "comp": self._comp_storage_bytes,
        }

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ):
        if encoder_hidden_states is not None:
            raise NotImplementedError("HGRN2OuterBlock does not support cross-attention.")

        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim

        q = self.q_proj(hidden_states).view(B, L, H, d).transpose(1, 2).contiguous()
        k = self.k_proj(hidden_states).view(B, L, self.num_kv_heads, d).transpose(1, 2).contiguous()
        v = self.v_proj(hidden_states).view(B, L, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # ---- RoPE ----
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
            cos, sin = pos_emb
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # ---- GQA ----
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.num_key_value_groups)
            v = repeat_kv(v, self.num_key_value_groups)

        # ---- Non-negative kernel ----
        q = F.elu(q) + 1.0
        k = (F.elu(k) + 1.0) / math.sqrt(d)

        # ---- Bounded forget gate ----
        f = torch.exp((-F.softplus(self.W_f(hidden_states))).clamp(min=-3.0, max=-1e-4)).transpose(
            1, 2
        )  # (B, H, L)

        # ---- Mask extraction (HGRN2-FIX-10) ----
        mask_bool = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                if attention_mask.dtype == torch.bool:
                    mask_bool = attention_mask
                elif attention_mask.dtype.is_floating_point:
                    mask_bool = attention_mask > -1e4
                else:
                    mask_bool = attention_mask > 0
            elif attention_mask.dim() == 4:
                if attention_mask.dtype == torch.bool:
                    m = attention_mask
                elif attention_mask.dtype.is_floating_point:
                    m = attention_mask > -1e4
                else:
                    m = attention_mask > 0
                mask_bool = m.any(dim=-2).squeeze(1)

            if mask_bool is not None and mask_bool.shape[-1] != L:
                mask_bool = None

        # Pre-mask k, v
        if mask_bool is not None:
            m_exp = mask_bool.unsqueeze(1).unsqueeze(-1).to(k.dtype)  # (B, 1, L, 1)
            k = k * m_exp
            v = v * m_exp

        # ---- Chunked scan ----
        y_fwd = self._chunked_scan(q, k, v, f, chunk_size=self.chunk_size)

        if self.is_decoder:
            y = y_fwd
        else:
            q_r = torch.flip(q, dims=[2])
            k_r = torch.flip(k, dims=[2])
            v_r = torch.flip(v, dims=[2])
            f_r = torch.flip(f, dims=[2])
            y_bwd_r = self._chunked_scan(q_r, k_r, v_r, f_r, chunk_size=self.chunk_size)
            y_bwd = torch.flip(y_bwd_r, dims=[2])
            y = 0.5 * (y_fwd + y_bwd)
            y = y.to(hidden_states.dtype)

        sn_dtype = (
            self.state_norm.weight.dtype
            if hasattr(self.state_norm, "weight") and self.state_norm.weight is not None
            else hidden_states.dtype
        )
        y = self.state_norm(y.to(sn_dtype)).to(hidden_states.dtype)
        y = torch.nan_to_num(y, nan=0.0, posinf=1e4, neginf=-1e4)

        # ---- head_mask ----
        if head_mask is not None:
            hm = head_mask.to(y.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(-1)
            y = y * hm

        # ---- output ----
        y = y.transpose(1, 2).contiguous().view(B, L, H * d)

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

        if mask_bool is not None:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        return out, None

Every 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_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.

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 storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    return {
        "orig": self._orig_storage_bytes,
        "comp": self._comp_storage_bytes,
    }

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class HLAAttention (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
num_heads: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("hla", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("hla_attention", supports_causal=True, supports_gqa=True)
class HLAAttention(BaseAttention):
    r"""Hadamard Linear Attention (HLA) Bidirectional / Causal Module.

    Maintains multi-head associative causal/anti-causal matrix states $S_t \in \mathbb{R}^{B \times H \times d_k \times d_k}$.
    """

    supports_gqa: bool = True
    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        if d_model is None:
            if config is not None:
                d_model = getattr(config, "hidden_size", getattr(config, "dim", 768))
            else:
                d_model = kwargs.get("hidden_size", 768)

        if hasattr(d_model, "hidden_size"):
            d_model = d_model.hidden_size

        if num_heads is None:
            if config is not None:
                num_heads = getattr(config, "num_attention_heads", getattr(config, "n_heads", 12))
            else:
                num_heads = kwargs.get("num_attention_heads", 12)

        comp_cfg = kwargs.get("comp_config")
        if comp_cfg is not None:
            lora_rank = getattr(comp_cfg, "lora_rank", getattr(comp_cfg, "rank", lora_rank))
            lora_alpha = getattr(comp_cfg, "lora_alpha", getattr(comp_cfg, "alpha", lora_alpha))
        elif config is not None:
            lora_rank = getattr(config, "lora_rank", lora_rank)
            lora_alpha = getattr(config, "lora_alpha", lora_alpha)

        self.hidden_size: int = int(d_model)
        self.num_heads: int = int(num_heads)
        self.head_dim: int = self.hidden_size // self.num_heads if self.num_heads > 0 else 64
        self.lora_rank: int = int(lora_rank) if lora_rank is not None else 16
        self.lora_alpha: float = float(lora_alpha) if lora_alpha is not None else 32.0

        cfg_kv = (
            getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
            if config is not None
            else None
        )
        self.num_kv_heads: int = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_key_value_groups: int = max(1, self.num_heads // max(1, self.num_kv_heads))

        self.q_head_size: int = self.num_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        # Decoder / Causal flag
        if is_decoder is not None:
            self.is_decoder: bool = is_decoder
        elif config is not None:
            self.is_decoder = bool(getattr(config, "is_decoder", False))
        else:
            self.is_decoder = False

        self.is_nested_encoder: bool = False

        self.q_proj = LoRALinear(
            self.hidden_size,
            self.q_head_size,
            bias=False,
            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=False,
            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=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.out_proj = LoRALinear(
            self.q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        self.beta = nn.Parameter(torch.ones(self.num_heads, 1, 1) * 2.0)
        self.output_scale = nn.Parameter(torch.ones(1))  # trainable output calibration

        old_attn = kwargs.get("old_attn", kwargs.get("pretrained_attn"))
        if old_attn is not None:
            self.copy_weights(old_attn)

    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Check nested encoder container
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                (src.c_attn.bias.detach() if src.c_attn.bias is not None else None)
                if c_attn_w.shape[0] == self.hidden_size and c_attn_w.shape[1] != self.hidden_size:
                    c_attn_w = c_attn_w.t().contiguous()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size

                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                out_w = src.c_proj.weight.detach()
                if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None

                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None

                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # Pattern 5: UniversalAdapter fallback
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    w_q, b_q = adapter.get_weight_bias(src, "q")
                    w_k, b_k = adapter.get_weight_bias(src, "k")
                    w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # Strict validation
            missing = []
            if w_q is None:
                missing.append("Q projection")
            if w_k is None:
                missing.append("K projection")
            if w_v is None:
                missing.append("V projection")
            if w_o is None and not self.is_nested_encoder:
                missing.append("Out projection")

            if missing:
                raise RuntimeError(
                    f"HLAAttention: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
                )

            for target_mod, w_src in [
                (self.q_proj, w_q),
                (self.k_proj, w_k),
                (self.v_proj, w_v),
            ]:
                tw = target_mod.weight
                w_cast = w_src.to(tw.dtype)
                if w_cast.shape == tw.shape:
                    tw.data.copy_(w_cast)
                else:
                    min_r = min(tw.shape[0], w_cast.shape[0])
                    min_c = min(tw.shape[1], w_cast.shape[1])
                    tw.data[:min_r, :min_c].copy_(w_cast[:min_r, :min_c])

            f_norm_q = float(torch.linalg.norm(self.q_proj.weight))
            if f_norm_q < 1e-4:
                logger.warning(f"HLAAttention: Q weight norm is near zero: {f_norm_q}")

            if w_o is not None:
                tw_o = self.out_proj.weight
                w_o_cast = w_o.to(tw_o.dtype)
                if w_o_cast.shape == tw_o.shape:
                    tw_o.data.copy_(w_o_cast)
                else:
                    min_r = min(tw_o.shape[0], w_o_cast.shape[0])
                    min_c = min(tw_o.shape[1], w_o_cast.shape[1])
                    tw_o.data[:min_r, :min_c].copy_(w_o_cast[:min_r, :min_c])

                if b_o is not None and self.out_proj.bias is not None:
                    tb_o = self.out_proj.bias
                    b_o_cast = b_o.to(tb_o.dtype)
                    min_b = min(tb_o.numel(), b_o_cast.numel())
                    tb_o.data[:min_b].copy_(b_o_cast[:min_b])
                f_norm_o = float(torch.linalg.norm(self.out_proj.weight))
                if f_norm_o < 1e-4:
                    logger.warning(f"HLAAttention: Out projection norm is near zero: {f_norm_o}")

    def _hla_scan(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        decay: torch.Tensor,
        mask: torch.Tensor | None = None,
        reverse: bool = False,
    ) -> torch.Tensor:
        """Vectorized parallel linear associative scan via tensor contraction."""
        B, H, L, d = q.shape
        q_f = torch.nn.functional.elu(q.float()) + 1.0
        k_f = torch.nn.functional.elu(k.float()) + 1.0
        v_f = v.float()

        if mask is not None:
            m_f = mask.float().view(B, 1, L, 1)
            k_f = k_f * m_f
            v_f = v_f * m_f

        t_idx = torch.arange(L, device=q.device)
        diff = (
            (t_idx.unsqueeze(0) - t_idx.unsqueeze(1))
            if reverse
            else (t_idx.unsqueeze(1) - t_idx.unsqueeze(0))
        )
        tri_mask = (
            torch.triu(torch.ones(L, L, device=q.device, dtype=torch.bool))
            if reverse
            else torch.tril(torch.ones(L, L, device=q.device, dtype=torch.bool))
        )

        log_decay = torch.log(
            decay.squeeze(-1).squeeze(-1).clamp(min=1e-5, max=1.0 - 1e-5)
        )  # [1, H]
        exp_arg = torch.where(
            tri_mask.unsqueeze(0).unsqueeze(0),
            diff.unsqueeze(0).unsqueeze(0).float() * log_decay.unsqueeze(-1).unsqueeze(-1),
            float("-inf"),
        )
        decay_mat = torch.exp(exp_arg)  # [1, H, L, L], exactly 0.0 outside causal triangle

        # Attention weights with exponential decay: scores = (q @ k^T) * decay_mat
        qk = torch.matmul(q_f, k_f.transpose(-1, -2).contiguous())  # [B, H, L, L]
        scores = qk * decay_mat

        # Numerator: y_num = scores @ v
        num = torch.matmul(scores, v_f)  # [B, H, L, d]

        # Denominator: y_denom = sum_j scores
        denom = scores.sum(dim=-1, keepdim=True) + 1e-5  # [B, H, L, 1]

        y = num / denom.clamp(min=1e-5)
        return y.to(q.dtype)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        """Bidirectional / Causal linear attention forward pass."""
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, L, D = hidden_states.shape
        H, Dh = self.num_heads, self.head_dim

        q = self.q_proj(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
        k = (
            self.k_proj(hidden_states)
            .view(B, L, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, L, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )

        # Partial Rotary Position Embeddings (RoPE)
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
            cos, sin = pos_emb
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # Grouped-Query Attention expansion
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.num_key_value_groups)
            v = repeat_kv(v, self.num_key_value_groups)

        # ELU + 1 kernelization
        q = F.elu(q) + 1.0
        k = F.elu(k) + 1.0

        decay = torch.sigmoid(self.beta).unsqueeze(0)  # [1, H, 1, 1]

        # Mask processing
        mask_bool = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                mask_bool = attention_mask.bool()
            elif attention_mask.dim() == 4:
                mask_bool = (attention_mask[:, 0, 0, :] > -10000.0).bool()

        y_fwd = self._hla_scan(q, k, v, decay, mask=mask_bool, reverse=False)

        if self.is_decoder:
            y = y_fwd
        else:
            y_bwd = self._hla_scan(q, k, v, decay, mask=mask_bool, reverse=True)
            y = 0.5 * (y_fwd + y_bwd)

        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)
            y = y * head_mask

        context = y.transpose(1, 2).contiguous().view(B, L, H * Dh)

        if self.is_nested_encoder:
            out = context * self.output_scale
        else:
            out = self.out_proj(context) * self.output_scale

        if attention_mask is not None and mask_bool is not None:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        return (out, None)

Hadamard Linear Attention (HLA) Bidirectional / Causal Module.

Maintains multi-head associative causal/anti-causal matrix states $S_t \in \mathbb{R}^{B \times H \times d_k \times d_k}$.

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 copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Check nested encoder container
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight
                b_o = old_attn.output.dense.bias
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            (src.c_attn.bias.detach() if src.c_attn.bias is not None else None)
            if c_attn_w.shape[0] == self.hidden_size and c_attn_w.shape[1] != self.hidden_size:
                c_attn_w = c_attn_w.t().contiguous()

            q_dim = self.q_head_size
            k_dim = self.kv_head_size
            v_dim = self.kv_head_size

            w_q = c_attn_w[:q_dim, :]
            w_k = c_attn_w[q_dim : q_dim + k_dim, :]
            w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

            out_w = src.c_proj.weight.detach()
            if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                out_w = out_w.t().contiguous()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None

            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None

            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # Pattern 5: UniversalAdapter fallback
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                w_q, b_q = adapter.get_weight_bias(src, "q")
                w_k, b_k = adapter.get_weight_bias(src, "k")
                w_v, b_v = adapter.get_weight_bias(src, "v")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(src, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # Strict validation
        missing = []
        if w_q is None:
            missing.append("Q projection")
        if w_k is None:
            missing.append("K projection")
        if w_v is None:
            missing.append("V projection")
        if w_o is None and not self.is_nested_encoder:
            missing.append("Out projection")

        if missing:
            raise RuntimeError(
                f"HLAAttention: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
            )

        for target_mod, w_src in [
            (self.q_proj, w_q),
            (self.k_proj, w_k),
            (self.v_proj, w_v),
        ]:
            tw = target_mod.weight
            w_cast = w_src.to(tw.dtype)
            if w_cast.shape == tw.shape:
                tw.data.copy_(w_cast)
            else:
                min_r = min(tw.shape[0], w_cast.shape[0])
                min_c = min(tw.shape[1], w_cast.shape[1])
                tw.data[:min_r, :min_c].copy_(w_cast[:min_r, :min_c])

        f_norm_q = float(torch.linalg.norm(self.q_proj.weight))
        if f_norm_q < 1e-4:
            logger.warning(f"HLAAttention: Q weight norm is near zero: {f_norm_q}")

        if w_o is not None:
            tw_o = self.out_proj.weight
            w_o_cast = w_o.to(tw_o.dtype)
            if w_o_cast.shape == tw_o.shape:
                tw_o.data.copy_(w_o_cast)
            else:
                min_r = min(tw_o.shape[0], w_o_cast.shape[0])
                min_c = min(tw_o.shape[1], w_o_cast.shape[1])
                tw_o.data[:min_r, :min_c].copy_(w_o_cast[:min_r, :min_c])

            if b_o is not None and self.out_proj.bias is not None:
                tb_o = self.out_proj.bias
                b_o_cast = b_o.to(tb_o.dtype)
                min_b = min(tb_o.numel(), b_o_cast.numel())
                tb_o.data[:min_b].copy_(b_o_cast[:min_b])
            f_norm_o = float(torch.linalg.norm(self.out_proj.weight))
            if f_norm_o < 1e-4:
                logger.warning(f"HLAAttention: Out projection norm is near zero: {f_norm_o}")

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

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

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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
    """Bidirectional / Causal linear attention forward pass."""
    if not hidden_states.is_contiguous():
        hidden_states = hidden_states.contiguous()

    B, L, D = hidden_states.shape
    H, Dh = self.num_heads, self.head_dim

    q = self.q_proj(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
    k = (
        self.k_proj(hidden_states)
        .view(B, L, self.num_kv_heads, Dh)
        .transpose(1, 2)
        .contiguous()
    )
    v = (
        self.v_proj(hidden_states)
        .view(B, L, self.num_kv_heads, Dh)
        .transpose(1, 2)
        .contiguous()
    )

    # Partial Rotary Position Embeddings (RoPE)
    pos_emb = kwargs.get("position_embeddings")
    if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
        cos, sin = pos_emb
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    # Grouped-Query Attention expansion
    if self.num_kv_heads != self.num_heads:
        k = repeat_kv(k, self.num_key_value_groups)
        v = repeat_kv(v, self.num_key_value_groups)

    # ELU + 1 kernelization
    q = F.elu(q) + 1.0
    k = F.elu(k) + 1.0

    decay = torch.sigmoid(self.beta).unsqueeze(0)  # [1, H, 1, 1]

    # Mask processing
    mask_bool = None
    if attention_mask is not None:
        if attention_mask.dim() == 2:
            mask_bool = attention_mask.bool()
        elif attention_mask.dim() == 4:
            mask_bool = (attention_mask[:, 0, 0, :] > -10000.0).bool()

    y_fwd = self._hla_scan(q, k, v, decay, mask=mask_bool, reverse=False)

    if self.is_decoder:
        y = y_fwd
    else:
        y_bwd = self._hla_scan(q, k, v, decay, mask=mask_bool, reverse=True)
        y = 0.5 * (y_fwd + y_bwd)

    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)
        y = y * head_mask

    context = y.transpose(1, 2).contiguous().view(B, L, H * Dh)

    if self.is_nested_encoder:
        out = context * self.output_scale
    else:
        out = self.out_proj(context) * self.output_scale

    if attention_mask is not None and mask_bool is not None:
        out = out * mask_bool.unsqueeze(-1).to(out.dtype)

    return (out, None)

Bidirectional / Causal linear attention forward pass.

Inherited members

• BaseAttention:
  • extra_repr

class HammingAttention (config: Any | None = None,
rank: int | None = None,
distill_temperature: float = 2.0,
sparsity_ratio: float = 0.0,
lora_rank: int = 16,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("hamming", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("hamming_attention", supports_causal=True, supports_gqa=True)
class HammingAttention(BaseAttention):
    supports_causal: 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,
        distill_temperature: float = 2.0,
        sparsity_ratio: float = 0.0,
        lora_rank: int = 16,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        num_heads, num_kv_heads, hidden_size, head_dim, kv_repeat = 0, 0, 0, 0, 1

        # HAD-FIX-07: ModelAdapter failures are surfaced (debug), not silently swallowed
        if config is not None and _HAS_MODEL_ADAPTER:
            try:
                self.adapter = ModelAdapter(config)
                num_heads = self.adapter.num_heads
                num_kv_heads = self.adapter.num_key_value_heads
                hidden_size = self.adapter.hidden_size
                head_dim = self.adapter.head_dim
                kv_repeat = self.adapter.kv_repeat
            except Exception as e:
                logger.debug("ModelAdapter failed: %s", e)

        # Fallback to direct config attributes
        if not hidden_size and config is not None:
            hidden_size = getattr(
                config,
                "hidden_size",
                getattr(config, "d_model", getattr(config, "dim", 0)),
            )
        if not num_heads and config is not None:
            num_heads = getattr(
                config,
                "num_attention_heads",
                getattr(config, "num_heads", getattr(config, "n_heads", 0)),
            )

        # Fallback to source attention module
        if not hidden_size and pretrained_attn is not None:
            hidden_size = getattr(
                pretrained_attn, "hidden_size", getattr(pretrained_attn, "embed_dim", 0)
            )
            if not hidden_size:
                for cand in ("q_proj", "q_lin", "query", "c_attn"):
                    m = getattr(pretrained_attn, cand, None)
                    if m is not None and hasattr(m, "in_features"):
                        hidden_size = m.in_features
                        break
                if not hidden_size:
                    out_l = getattr(
                        pretrained_attn,
                        "out_lin",
                        getattr(
                            pretrained_attn,
                            "dense",
                            getattr(pretrained_attn, "c_proj", None),
                        ),
                    )
                    if out_l is not None and hasattr(out_l, "out_features"):
                        hidden_size = out_l.out_features

        # Final defaults
        hidden_size = hidden_size or kwargs.get("hidden_size", kwargs.get("d_model", 768))
        num_heads = num_heads or kwargs.get("num_heads", kwargs.get("num_attention_heads", 12))
        num_kv_heads = num_kv_heads or _cfg_get(
            config, "num_key_value_heads", kwargs.get("num_key_value_heads", num_heads)
        )
        head_dim = head_dim or kwargs.get("head_dim") or (hidden_size // max(1, num_heads))
        kv_repeat = kv_repeat or max(1, num_heads // max(1, num_kv_heads))

        self.num_heads: int = int(num_heads)
        self.num_kv_heads: int = int(num_kv_heads)
        self.hidden_size: int = int(hidden_size)
        self.head_dim: int = int(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 = int(kv_repeat)

        # HAD-FIX-06: keep API compat args but they're informational
        self.distill_temperature: float = float(distill_temperature)
        self.sparsity_ratio: float = float(sparsity_ratio)

        # ---- Decoder detection ----
        model_type = str(_cfg_get(config, "model_type", "")).lower()
        is_dec = _cfg_get(
            config,
            "is_decoder",
            _cfg_get(
                config,
                "is_causal",
                kwargs.get("is_causal", kwargs.get("is_decoder", False)),
            ),
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if any(_mt.startswith(p) for p in _decoder_prefixes):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = bool(is_dec)
        self.is_nested_encoder: bool = False

        self.lora_rank = int(_cfg_get(config, "lora_rank", kwargs.get("lora_rank", lora_rank)))

        # ---- Projections ----
        self.q_proj = LoRALinear(
            self.hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.k_proj = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.v_proj = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )

        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)

        # HAD-FIX-01: rename trainable scalars with `lora_` prefix so that
        # framework's get_recovery_params recognizes them.
        self.lora_q_scale = nn.Parameter(torch.ones(self.num_heads, 1))
        self.lora_k_scale = nn.Parameter(torch.ones(self.num_heads, 1))
        self.lora_temp = nn.Parameter(torch.tensor(max(distill_temperature, 1e-2)))

        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

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

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. UniversalAdapter
        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
        if hasattr(src, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            weight, bias_data = _unpack_weight_bias(src.c_attn)
            if weight is not None:
                if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                    weight = weight.t().contiguous()
                split_size = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split_size, :],
                    bias_data[:split_size] if bias_data is not None else None,
                )
                extracted["k"] = (
                    weight[split_size : 2 * split_size, :],
                    (bias_data[split_size : 2 * split_size] if bias_data is not None else None),
                )
                extracted["v"] = (
                    weight[2 * split_size :, :],
                    bias_data[2 * split_size :] if bias_data is not None else None,
                )

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            w_o, b_o = _unpack_weight_bias(src.c_proj)
            if w_o is not None:
                if (
                    w_o.dim() == 2
                    and w_o.shape[0] != self.hidden_size
                    and w_o.shape[1] == self.hidden_size
                ):
                    w_o = w_o.t().contiguous()
                extracted["o"] = (w_o, b_o)

        # 3. RoBERTa / BERT nested
        if hasattr(src, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            s = src.self
            if hasattr(s, "query") and extracted["q"][0] is None:
                extracted["q"] = _unpack_weight_bias(s.query)
            if hasattr(s, "key") and extracted["k"][0] is None:
                extracted["k"] = _unpack_weight_bias(s.key)
            if hasattr(s, "value") and extracted["v"][0] is None:
                extracted["v"] = _unpack_weight_bias(s.value)

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            extracted["o"] = _unpack_weight_bias(src.output.dense)

        # 4. Flat modules
        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(src, name):
                        w, b = _unpack_weight_bias(getattr(src, name))
                        if w is not None:
                            extracted[proj] = (w, b)
                            break

        # ---- Nested encoder detection ----
        cls_name = getattr(src, "__class__", type(src)).__name__
        has_out = any(hasattr(src, a) for a in ("dense", "out_proj", "out_lin", "c_proj", "o_proj"))
        if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
            self.is_nested_encoder = False
        elif (
            (
                cls_name in ("BertSelfAttention", "RobertaSelfAttention")
                or (hasattr(src, "query") and not has_out)
            )
            or hasattr(src, "self")
            and hasattr(src, "output")
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        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.debug("HammingAttention partial extraction; missing: %s", missing)

        # Frobenius check
        for proj_key in ("q", "k", "v"):
            w = extracted[proj_key][0]
            if w is not None:
                n = float(torch.linalg.norm(w.float()).item())
                if n < 1e-4:
                    logger.debug("HammingAttention: %s near-zero norm (%.2e)", proj_key, n)
        if extracted["o"][0] is not None:
            n = float(torch.linalg.norm(extracted["o"][0].float()).item())
            if n < 1e-4:
                logger.debug("HammingAttention: out near-zero norm (%.2e)", n)

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

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

                target.freeze_base_weights()

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

    # ------------------------------------------------------------------
    # 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,
        *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, L_q, D = hidden_states.shape
        L_k = kv_input.shape[1]
        H, Dh = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = self.q_proj(hidden_states).view(B, L_q, H, Dh).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, L_k, self.num_kv_heads, Dh).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, L_k, self.num_kv_heads, Dh).transpose(1, 2).contiguous()

        # ---- RoPE ----
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None and not is_cross_attention:
            if 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"))

        # ---- GQA ----
        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        # ---- Hamming attention scores ----
        # Continuous relaxation of Hamming distance via tanh binarization.
        # Soft-binary: q_b, k_b ∈ (-1, 1)^Dh.
        # Inner product <q_b, k_b> ∈ (-Dh, Dh), so (matmul - Dh) is in (-2Dh, 0).
        # Dividing by 2*sqrt(Dh) normalizes; multiplying by learned temp sharpens.
        q_b = torch.tanh(5.0 * Q * self.lora_q_scale.view(1, H, 1, 1))
        k_b = torch.tanh(5.0 * K * self.lora_k_scale.view(1, H, 1, 1))

        # HAD-FIX-02: ensure temperature is positive (softmax over -Hamming needs T > 0)
        temp = F.softplus(self.lora_temp) + 1e-3

        scores = (
            temp
            * (torch.matmul(q_b, k_b.transpose(-2, -1).contiguous()) - Dh)
            / (2.0 * math.sqrt(Dh))
        )

        # ---- Masking ----
        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, L_k, device)
                if valid_m is not None:
                    scores = scores.masked_fill((~valid_m).unsqueeze(1).unsqueeze(2), -1e4)

        # ---- Causal mask ----
        if self.is_decoder and not is_cross_attention and L_q > 1:
            causal_tril = (
                torch.tril(
                    torch.ones((L_q, L_k), device=device, dtype=torch.bool),
                    diagonal=L_k - L_q,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores.masked_fill(~causal_tril, -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores.float(), dim=-1).to(dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        # ---- head_mask ----
        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn = attn * head_mask

        context = torch.matmul(attn, V)
        context = context.transpose(1, 2).contiguous().view(B, L_q, self.all_head_size)

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

        if output.shape[-1] != hidden_states.shape[-1]:
            if output.shape[-1] < hidden_states.shape[-1]:
                output = F.pad(output, (0, hidden_states.shape[-1] - output.shape[-1]))
            else:
                output = output[..., : hidden_states.shape[-1]]

        return output, (attn if output_attentions else None)

Every 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.

var supports_rope : bool

The type of the None singleton.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class HeadGate (hidden_dim: int, n_heads: int)

Expand source code

class HeadGate(nn.Module):
    r"""Learned per-example head importance scalars with masked pooling."""

    def __init__(self, hidden_dim: int, n_heads: int) -> None:
        super().__init__()
        self.net = nn.Sequential(nn.Linear(hidden_dim, 32), nn.GELU(), nn.Linear(32, n_heads))
        # Initialize head gate to neutral 1.0 (bias=3.0 -> sigmoid(3.0) ≈ 0.95)
        nn.init.zeros_(self.net[-1].weight)
        nn.init.constant_(self.net[-1].bias, 3.0)

    def forward(self, x: torch.Tensor, valid_mask: torch.Tensor | None = None) -> torch.Tensor:
        net_dtype = next(self.net.parameters()).dtype
        pooled = _masked_mean(x, valid_mask).to(net_dtype)
        return torch.sigmoid(self.net(pooled)).to(x.dtype)

Learned per-example head importance scalars with masked pooling.

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor, valid_mask: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor, valid_mask: torch.Tensor | None = None) -> torch.Tensor:
    net_dtype = next(self.net.parameters()).dtype
    pooled = _masked_mean(x, valid_mask).to(net_dtype)
    return torch.sigmoid(self.net(pooled)).to(x.dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class HeadRoPE (head_dim: int, max_seq_len: int = 4096, base: float = 10000.0)

Expand source code

class HeadRoPE(nn.Module):
    def __init__(self, head_dim: int, max_seq_len: int = 4096, base: float = 10000.0) -> None:
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._build_cache(max_seq_len)

    def _build_cache(self, seq_len: int) -> None:
        t = torch.arange(seq_len, device=self.inv_freq.device).float()
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().unsqueeze(0).unsqueeze(1), persistent=False)
        self.register_buffer("sin_cached", emb.sin().unsqueeze(0).unsqueeze(1), persistent=False)

    def forward(self, x: torch.Tensor, past_len: int = 0) -> torch.Tensor:
        B, H, L, d = x.shape
        if past_len + L > self.max_seq_len:
            self._build_cache(past_len + L)
            self.max_seq_len = past_len + L

        cos = self.cos_cached[:, :, past_len : past_len + L, :d]
        sin = self.sin_cached[:, :, past_len : past_len + L, :d]

        x_even = x[..., 0::2]
        x_odd = x[..., 1::2]
        rot_even = x_even * cos[..., 0::2] - x_odd * sin[..., 0::2]
        rot_odd = x_even * sin[..., 0::2] + x_odd * cos[..., 0::2]
        return torch.stack((rot_even, rot_odd), dim=-1).flatten(-2)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor, past_len: int = 0) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor, past_len: int = 0) -> torch.Tensor:
    B, H, L, d = x.shape
    if past_len + L > self.max_seq_len:
        self._build_cache(past_len + L)
        self.max_seq_len = past_len + L

    cos = self.cos_cached[:, :, past_len : past_len + L, :d]
    sin = self.sin_cached[:, :, past_len : past_len + L, :d]

    x_even = x[..., 0::2]
    x_odd = x[..., 1::2]
    rot_even = x_even * cos[..., 0::2] - x_odd * sin[..., 0::2]
    rot_odd = x_even * sin[..., 0::2] + x_odd * cos[..., 0::2]
    return torch.stack((rot_even, rot_odd), dim=-1).flatten(-2)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

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 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

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

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class HymbaTransXBlock (config: Any | None = None,
rank: int | None = 64,
num_heads: int | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "hymba_transx", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "hymba_transx_attn",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class HymbaTransXBlock(BaseAttention):
    r"""HymbaTransX: Hybrid Differential Transformer & State Space Model Block."""

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        num_heads: int | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        hidden_size = kwargs.get("hidden_size")
        num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")
        if hidden_size is None and config is not None:
            hidden_size = getattr(
                config,
                "hidden_size",
                getattr(
                    config,
                    "dim",
                    getattr(config, "d_model", getattr(config, "n_embd", None)),
                ),
            )
        if hidden_size is None and source_attn is not None:
            for attr in ["hidden_size", "embed_dim", "dim", "d_model"]:
                if hasattr(source_attn, attr) and getattr(source_attn, attr) is not None:
                    hidden_size = getattr(source_attn, attr)
                    break
            if hidden_size is None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                hidden_size = (
                    getattr(out_layer, "out_features", None) if out_layer is not None else None
                )
        hidden_size = int(hidden_size or 768)

        if num_heads is None and config is not None:
            num_heads = getattr(
                config,
                "num_attention_heads",
                getattr(config, "n_heads", getattr(config, "num_heads", None)),
            )
        if num_heads is None and source_attn is not None:
            num_heads = getattr(
                source_attn,
                "num_attention_heads",
                getattr(source_attn, "n_heads", getattr(source_attn, "num_heads", None)),
            )
        num_heads = int(num_heads or (hidden_size // 64 if hidden_size % 64 == 0 else 12))

        self.dim = hidden_size
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        cfg_kv = (
            getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
            if config
            else None
        )
        self.num_kv_heads = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))

        cfg_head_dim = (
            getattr(config, "head_dim", kwargs.get("head_dim"))
            if config
            else kwargs.get("head_dim")
        )
        if cfg_head_dim is None and source_attn is not None:
            for attr in ("q_proj", "query", "q_lin", "W_q1"):
                mod = getattr(source_attn, attr, None)
                if mod is not None and hasattr(mod, "weight"):
                    cfg_head_dim = mod.weight.shape[0] // max(1, self.num_heads)
                    break
        self.head_dim = (
            int(cfg_head_dim) if cfg_head_dim is not None else self.dim // max(1, self.num_heads)
        )
        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 // self.num_kv_heads)
        self.supports_gqa = True

        # Decoder detection
        self.is_decoder = _detect_decoder(config, is_decoder, source_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False
        # HX-FIX-02: RoPE only for models that expect it
        self.use_rope = _uses_rope(config, source_attn) or self.is_decoder

        # LoRA config
        self.lora_rank = int(
            kwargs.get(
                "lora_rank",
                getattr(
                    comp_config,
                    "lora_rank",
                    getattr(config, "lora_rank", 16) if config else 16,
                ),
            )
        )
        self.lora_alpha = float(
            kwargs.get(
                "lora_alpha",
                getattr(
                    comp_config,
                    "lora_alpha",
                    getattr(config, "lora_alpha", 32.0) if config else 32.0,
                ),
            )
        )
        self.freeze_base = self.lora_rank > 0

        # RoPE and projections
        self.rope = HeadRoPE(self.head_dim) if self.use_rope else None

        self.W_q1 = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_k1 = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_q2 = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_k2 = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_v = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # HX-FIX-06: W_q2 and W_k2 are noisy copies of W_q1 and W_k1.
        with torch.no_grad():
            noise_q = torch.randn_like(self.W_q2.weight) * 0.01
            noise_k = torch.randn_like(self.W_k2.weight) * 0.01
            self.W_q2.weight.data.copy_(self.W_q1.weight.data + noise_q)
            self.W_k2.weight.data.copy_(self.W_k1.weight.data + noise_k)

        # HX-FIX-07: differential attention lambda is learnable.
        self.diff_lambda = nn.Parameter(torch.tensor(0.2))

        # SSM branch
        try:
            from neuralops.attentions.simple_ssm import SimpleSSM

            self.mamba = SimpleSSM(d_model=self.dim, config=config, is_decoder=self.is_decoder)
        except ImportError:
            self.mamba = None
            logger.warning("HymbaTransX: SimpleSSM not available; SSM branch disabled.")

        # HX-FIX-01: alpha init = +1.0 (attention-dominant)
        self.alpha = nn.Parameter(torch.tensor(1.0))

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

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

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

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

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

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

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

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

    @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

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    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")
        if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
            self.is_nested_encoder = False
        elif (
            (
                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
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        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
        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()
                cls_c = c_attn_mod.__class__.__name__
                base_cls_c = base_c.__class__.__name__
                if (
                    "Conv1D" in cls_c
                    or "Conv1D" in base_cls_c
                    or (c_w.ndim == 2 and c_w.shape[1] == 3 * self.dim)
                ):
                    c_w = c_w.t().contiguous()
                split = 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, :],
                    c_b[:split] if c_b is not None else None,
                )
                extracted["k"] = (
                    c_w[split : 2 * split, :],
                    c_b[split : 2 * split] if c_b is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :, :],
                    c_b[2 * split :] 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()
                    cls_o = c_proj_mod.__class__.__name__
                    base_cls_o = base_o.__class__.__name__
                    if (
                        "Conv1D" in cls_o
                        or "Conv1D" in base_cls_o
                        or (o_w.ndim == 2 and o_w.shape[0] != self.dim and o_w.shape[1] == self.dim)
                    ):
                        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. RoBERTa / BERT nested
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                mod = getattr(s, sub, None)
                if mod is not None and hasattr(mod, "weight"):
                    extracted[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
        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 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[key] = (
                                mod.weight.data.clone(),
                                mod.bias.data.clone() if mod.bias is not None else None,
                            )
                            break

        # 4. UniversalAdapter fallback
        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)

        with torch.no_grad():
            w_q, _ = extracted["q"]
            w_k, _ = extracted["k"]
            w_v, _ = extracted["v"]
            w_o, b_o = extracted["o"]

            def _adapt(w, target_shape):
                if w is None:
                    return None
                if w.ndim != 2:
                    w = w.reshape(w.shape[0], -1)
                if w.shape == target_shape:
                    return w
                if w.t().contiguous().shape == target_shape:
                    return w.t().contiguous()
                # Truncate / pad
                out_d, in_d = target_shape
                s_out, s_in = w.shape
                m_out = min(out_d, s_out)
                m_in = min(in_d, s_in)
                padded = w.new_zeros(target_shape)
                padded[:m_out, :m_in] = w[:m_out, :m_in]
                return padded

            # W_q1
            wq = _adapt(w_q, self.W_q1.weight.shape)
            if wq is not None:
                self.W_q1.weight.data.copy_(
                    wq.to(device=self.W_q1.weight.device, dtype=self.W_q1.weight.dtype)
                )

            # W_k1
            wk = _adapt(w_k, self.W_k1.weight.shape)
            if wk is not None:
                self.W_k1.weight.data.copy_(
                    wk.to(device=self.W_k1.weight.device, dtype=self.W_k1.weight.dtype)
                )

            # W_v
            wv = _adapt(w_v, self.W_v.weight.shape)
            if wv is not None:
                self.W_v.weight.data.copy_(
                    wv.to(device=self.W_v.weight.device, dtype=self.W_v.weight.dtype)
                )

            # W_o
            if w_o is not None:
                wo = _adapt(w_o, self.out_proj.weight.shape)
                if wo is not None:
                    self.out_proj.weight.data.copy_(
                        wo.to(
                            device=self.out_proj.weight.device,
                            dtype=self.out_proj.weight.dtype,
                        )
                    )
                if b_o is not None and self.out_proj.bias is not None:
                    m = min(self.out_proj.bias.numel(), b_o.numel())
                    self.out_proj.bias.data.zero_()
                    self.out_proj.bias.data[:m].copy_(
                        b_o[:m].to(
                            device=self.out_proj.bias.device,
                            dtype=self.out_proj.bias.dtype,
                        )
                    )

            # HX-FIX-06: W_q2 and W_k2 = noisy copies of W_q1 and W_k1
            noise_q = torch.randn_like(self.W_q2.weight) * 0.01
            noise_k = torch.randn_like(self.W_k2.weight) * 0.01
            self.W_q2.weight.data.copy_(self.W_q1.weight.data + noise_q.to(self.W_q2.weight.dtype))
            self.W_k2.weight.data.copy_(self.W_k1.weight.data + noise_k.to(self.W_k2.weight.dtype))

    # ------------------------------------------------------------------
    # Mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # 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]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("HymbaTransXBlock does not support cross-attention.")

        x = hidden_states if hidden_states is not None else kwargs.get("query", kwargs.get("x"))
        if x is None and len(args) > 0:
            x = args[0]
        if x is None:
            raise ValueError("HymbaTransXBlock.forward: no input tensor provided.")

        B, L, D = x.shape
        device = x.device

        # Projections
        q1 = self.W_q1(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k1 = self.W_k1(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        q2 = self.W_q2(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k2 = self.W_k2(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        v = self.W_v(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()

        # HX-FIX-04: KV-cache.
        past_k1 = past_v = None
        past_k2 = None
        past_len = 0
        if past_key_value is not None:
            if isinstance(past_key_value, (tuple, list)):
                # Expect (k1, v, k2)
                if len(past_key_value) >= 3:
                    past_k1, past_v, past_k2 = (
                        past_key_value[0],
                        past_key_value[1],
                        past_key_value[2],
                    )
                elif len(past_key_value) >= 2:
                    past_k1, 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_k1 = past_key_value.key_cache[idx]
                past_v = past_key_value.value_cache[idx]

        if past_k1 is not None:
            past_len = past_k1.shape[2]
            k1_cat = torch.cat([past_k1, k1], dim=2)
            v_cat = torch.cat([past_v, v], dim=2)
            k2_cat = torch.cat([past_k2, k2], dim=2) if past_k2 is not None else k2
        else:
            k1_cat = k1
            v_cat = v
            k2_cat = k2

        N_kv = k1_cat.shape[2]

        # RoPE (HX-FIX-02: only if use_rope)
        if self.use_rope:
            if position_embeddings is not None:
                cos, sin = position_embeddings
                cos = cos[..., past_len : past_len + L, :]
                sin = sin[..., past_len : past_len + L, :]
                q1, k1_cat = apply_rotary_pos_emb(q1, k1_cat, cos, sin)
                q2, k2_cat = apply_rotary_pos_emb(q2, k2_cat, cos, sin)
            elif self.rope is not None:
                q1 = self.rope(q1, past_len=past_len)
                k1_cat = self.rope(k1_cat, past_len=0)
                q2 = self.rope(q2, past_len=past_len)
                k2_cat = self.rope(k2_cat, past_len=0)

        # GQA expansion
        if self.kv_repeat > 1:
            k1_cat = repeat_kv(k1_cat, self.kv_repeat)
            k2_cat = repeat_kv(k2_cat, self.kv_repeat)
            v_cat = repeat_kv(v_cat, self.kv_repeat)

        # Differential attention
        scale = 1.0 / math.sqrt(self.head_dim)
        scores1 = torch.matmul(q1.float(), k1_cat.float().transpose(-2, -1).contiguous()) * scale
        scores2 = torch.matmul(q2.float(), k2_cat.float().transpose(-2, -1).contiguous()) * scale

        # Causal mask (HX-FIX-03: offset-aware for cache)
        if self.is_decoder:
            if past_len > 0:
                abs_q = (past_len + torch.arange(L, device=device)).unsqueeze(1)
                abs_k = torch.arange(N_kv, device=device).unsqueeze(0)
                causal = abs_k <= abs_q
            else:
                causal = torch.tril(torch.ones(L, N_kv, device=device, dtype=torch.bool))
            scores1 = scores1.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)
            scores2 = scores2.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

        # Padding mask (HX-FIX-12)
        eff_mask = encoder_attention_mask if encoder_hidden_states is not None else attention_mask
        valid_mask = self._prepare_valid_mask(eff_mask, B, N_kv, device)
        if valid_mask is not None:
            scores1 = scores1.masked_fill(~valid_mask[:, None, None, :].bool(), -1e4)
            scores2 = scores2.masked_fill(~valid_mask[:, None, None, :].bool(), -1e4)

        scores1 = torch.nan_to_num(scores1, nan=-1e4, posinf=1e4, neginf=-1e4)
        scores2 = torch.nan_to_num(scores2, nan=-1e4, posinf=1e4, neginf=-1e4)

        attn1 = F.softmax(scores1, dim=-1)
        attn2 = F.softmax(scores2, dim=-1)

        # HX-FIX-07: learnable lambda
        lam = torch.sigmoid(self.diff_lambda)
        diff_attn = F.relu(attn1 - lam * attn2)
        denom = diff_attn.sum(dim=-1, keepdim=True)
        diff_attn = diff_attn / (denom + 1e-6)
        diff_attn = self.dropout(diff_attn).to(v_cat.dtype)

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

        A_out = (
            torch.matmul(diff_attn, v_cat)
            .transpose(1, 2)
            .contiguous()
            .view(B, L, self.all_head_size)
        )

        # SSM branch
        if self.mamba is not None:
            if x.requires_grad:
                from torch.utils.checkpoint import checkpoint

                def _run_mamba(h, m):
                    return self.mamba(h, attention_mask=m)[0]

                if self.is_decoder:
                    M_out = checkpoint(_run_mamba, x, attention_mask, use_reentrant=False)
                else:
                    x_rev = torch.flip(x, dims=[1])
                    mf = checkpoint(_run_mamba, x, attention_mask, use_reentrant=False)
                    mb = checkpoint(_run_mamba, x_rev, attention_mask, use_reentrant=False)
                    M_out = 0.5 * (mf + torch.flip(mb, dims=[1]))
            else:
                if self.is_decoder:
                    mamba_res = self.mamba(x, attention_mask=attention_mask)
                    M_out = mamba_res[0] if isinstance(mamba_res, tuple) else mamba_res
                else:
                    x_rev = torch.flip(x, dims=[1])
                    mamba_fwd = self.mamba(x, attention_mask=attention_mask)
                    mamba_bwd = self.mamba(x_rev, attention_mask=attention_mask)
                    mf = mamba_fwd[0] if isinstance(mamba_fwd, tuple) else mamba_fwd
                    mb = mamba_bwd[0] if isinstance(mamba_bwd, tuple) else mamba_bwd
                    M_out = 0.5 * (mf + torch.flip(mb, dims=[1]))
        else:
            M_out = A_out

        # HX-FIX-01: alpha init = +1.0
        alpha = torch.sigmoid(self.alpha)
        combined = alpha * A_out + (1.0 - alpha) * M_out

        # Output projection
        if self.is_nested_encoder:
            out = combined
        else:
            out = self.out_proj(combined)

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

HymbaTransX: Hybrid Differential Transformer & State Space Model Block.

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.

Instance variables

prop dense

Expand source code

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

prop k_proj

Expand source code

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

prop key

Expand source code

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

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_proj

Expand source code

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

prop query

Expand source code

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

prop v_proj

Expand source code

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

prop value

Expand source code

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

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class IdentityOutProj

Expand source code

class IdentityOutProj(nn.Identity):
    def __init__(self) -> None:
        super().__init__()
        self.lora_A = None
        self.lora_B = None

    def freeze_base_weights(self) -> None:
        pass

A placeholder identity operator that is argument-insensitive.

Args

-----=

args

any argument (unused)

kwargs

any keyword argument (unused)

Shape -----= - Input: :math:(*), where :math:* means any number of dimensions. - Output: :math:(*), same shape as the input.

Examples::

>>> m = nn.Identity(54, unused_argument1=0.1, unused_argument2=False)
>>> input = torch.randn(128, 20)
>>> output = m(input)
>>> print(output.size())
torch.Size([128, 20])

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

Ancestors

• torch.nn.modules.linear.Identity
• torch.nn.modules.module.Module

Methods

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    pass

class InfiniAttention (config: Any | None = None,
rank: int | None = None,
local_window: int = 128,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("infini", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("infini_attention", supports_causal=True, supports_gqa=True)
class InfiniAttention(BaseAttention):
    """
    Infini-Attention: sliding-window local attention + compressive memory.

    For decoders, the output is a per-head convex combination of:
      - A_local: sliding-window softmax attention
      - A_mem:   read from a transient (per-forward) compressive memory state
    For encoders, it degenerates to full softmax attention (memory is skipped).
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        local_window: int = 128,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | 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 if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        # ---------- dim ----------
        hidden_size = kwargs.get("hidden_size")
        num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")
        num_kv_heads = kwargs.get("num_key_value_heads") or kwargs.get("num_kv_heads")

        if hidden_size is None and config is not None:
            hidden_size = getattr(config, "hidden_size", getattr(config, "dim", 768))
        if hidden_size is None and source_attn is not None:
            for attr in ("out_lin", "dense", "c_proj", "o_proj", "out_proj"):
                layer = getattr(source_attn, attr, None)
                if layer is not None and hasattr(layer, "out_features"):
                    hidden_size = layer.out_features
                    break
        hidden_size = int(hidden_size or 768)

        if num_heads is None and config is not None:
            num_heads = getattr(
                config,
                "num_attention_heads",
                getattr(config, "n_heads", getattr(config, "num_heads", None)),
            )
        if num_heads is None and source_attn is not None:
            num_heads = getattr(
                source_attn,
                "num_attention_heads",
                getattr(source_attn, "n_heads", None),
            )
        num_heads = int(num_heads or 12)

        if num_kv_heads is None and config is not None:
            num_kv_heads = getattr(
                config,
                "num_key_value_heads",
                getattr(config, "num_kv_heads", num_heads),
            )
        if num_kv_heads is None:
            num_kv_heads = num_heads
        num_kv_heads = int(num_kv_heads)

        self.dim = hidden_size
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_kv_heads = max(1, min(num_kv_heads, num_heads))

        # ---------- head_dim ----------
        inferred = None
        if config is not None:
            inferred = getattr(config, "head_dim", None)
        if inferred is None:
            inferred = kwargs.get("head_dim")
        if inferred is None and source_attn is not None:
            for attr in ("q_proj", "query", "q_lin"):
                layer = getattr(source_attn, attr, None)
                if layer is not None and hasattr(layer, "weight"):
                    inferred = layer.weight.shape[0] // max(1, num_heads)
                    break
        self.head_dim = int(inferred or (self.dim // max(1, num_heads)))
        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

        # ---------- window ----------
        self.local_window = max(1, int(local_window))

        # ---------- decoder ----------
        self.is_decoder = _detect_decoder(config, is_decoder, source_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # ---------- LoRA ----------
        self.lora_rank = int(
            kwargs.get(
                "lora_rank",
                _cfg_get(
                    comp_config,
                    "lora_rank",
                    getattr(config, "lora_rank", 16) if config else 16,
                ),
            )
        )
        self.lora_alpha = float(
            kwargs.get(
                "lora_alpha",
                _cfg_get(
                    comp_config,
                    "lora_alpha",
                    getattr(config, "lora_alpha", 32.0) if config else 32.0,
                ),
            )
        )
        self.freeze_base = self.lora_rank > 0

        # ---------- projections ----------
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # ---------- fusion gate (per-head, init at sigmoid(0)=0.5) ----------
        self.beta = nn.Parameter(torch.zeros(self.num_heads))

        # ---------- write gate (per-token, input-dependent) ----------
        self.w_mem = nn.Linear(self.dim, self.num_heads, bias=True)
        with torch.no_grad():
            nn.init.zeros_(self.w_mem.weight)
            nn.init.constant_(self.w_mem.bias, 2.0)  # sigmoid(2) ≈ 0.88

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.0))
            if config
            else 0.0
        )
        self.dropout = nn.Dropout(dropout_p)

        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

        if source_attn is not None:
            self.copy_weights(source_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 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

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

        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        cls_name = src.__class__.__name__
        has_out = any(hasattr(src, a) for a in ("dense", "out_proj", "out_lin", "c_proj", "o_proj"))

        is_nested = False
        if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
            is_nested = False
        elif (
            (
                hasattr(src, "self")
                and hasattr(src, "output")
                and hasattr(src.self, "query")
                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
        ):
            is_nested = True
        self.is_nested_encoder = is_nested

        extracted = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c = src.c_attn
            base_c = getattr(c, "base_layer", c)
            cw = getattr(base_c, "weight", getattr(c, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                is_conv = (
                    "Conv1D" in c.__class__.__name__
                    or "Conv1D" in base_c.__class__.__name__
                    or hasattr(c, "nf")
                    or hasattr(base_c, "nf")
                    or (cw.ndim == 2 and cw.shape[1] == 3 * self.dim)
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )

            if hasattr(src, "c_proj") and src.c_proj is not None:
                o = src.c_proj
                base_o = getattr(o, "base_layer", o)
                ow = getattr(base_o, "weight", getattr(o, "weight", None))
                if ow is not None:
                    ow = ow.data.clone()
                    is_conv_o = (
                        "Conv1D" in o.__class__.__name__
                        or "Conv1D" in base_o.__class__.__name__
                        or hasattr(o, "nf")
                        or hasattr(base_o, "nf")
                        or (ow.ndim == 2 and ow.shape[0] != self.dim and ow.shape[1] == self.dim)
                    )
                    if is_conv_o:
                        ow = ow.t().contiguous()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # Nested BERT / RoBERTa
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                d = src.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat (DistilBERT, Qwen, LLaMA)
        else:
            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 key, names in candidates.items():
                for n in names:
                    if hasattr(src, n):
                        m = getattr(src, n)
                        if m is not None and hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data.clone(),
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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("Infini UniversalAdapter fallback failed: %s", e)

        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

        with torch.no_grad():
            for key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target.weight)
                        if target.bias is not None:
                            nn.init.zeros_(target.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target.weight.shape
                    and w.t().contiguous().shape == target.weight.shape
                ):
                    w = w.t().contiguous()

                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))
                target.weight.data.zero_()
                target.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                )
                if b is not None and target.bias is not None:
                    mb = min(target.bias.numel(), b.numel())
                    target.bias.data.zero_()
                    target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))
                elif target.bias is not None:
                    nn.init.zeros_(target.bias)

                n = torch.linalg.norm(target.weight.data.float()).item()
                if n < 1e-4:
                    logger.warning("Infini: '%s' near-zero norm (%.2e)", key, n)

                if self.freeze_base:
                    target.freeze_base_weights()

        # Storage accounting
        try:
            self._comp_storage_bytes = sum(
                p.numel() * p.element_size()
                for m in (
                    self.q_proj,
                    self.k_proj,
                    self.v_proj,
                    self.out_proj,
                    self.w_mem,
                )
                for p in m.parameters()
            )
        except Exception:
            pass

    # ------------------------------------------------------------------
    # Mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(attention_mask, B, L, device):
        if attention_mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape == (B, L):
                return m
        elif attention_mask.dim() == 3:
            m = _to_bool(attention_mask)
            if m.shape[0] == B and m.shape[-1] == L:
                return m[:, -1, :]
        elif attention_mask.dim() == 4:
            m = _to_bool(attention_mask)
            m = m.any(dim=-2)
            if m.dim() == 3:
                m = m.squeeze(1)
            if m.shape == (B, L):
                return m
        return None

    # ------------------------------------------------------------------
    # Sliding-window local attention
    # ------------------------------------------------------------------
    def _local_sliding_attention(self, q, k, v, valid_mask=None):
        B, H, L, Dh = q.shape
        scale = 1.0 / math.sqrt(Dh)

        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        idx = torch.arange(L, device=q.device)
        dist = (idx.unsqueeze(0) - idx.unsqueeze(1)).abs()
        band = dist <= self.local_window

        if self.is_decoder:
            causal = torch.tril(torch.ones(L, L, device=q.device, dtype=torch.bool))
            band = band & causal

        scores = scores.masked_fill(~band.unsqueeze(0).unsqueeze(0), -1e4)

        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(q.dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

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

    # ------------------------------------------------------------------
    # Full dense attention (encoder path)
    # ------------------------------------------------------------------
    def _full_attention(self, q, k, v, valid_mask=None):
        B, H, L, Dh = q.shape
        scale = 1.0 / math.sqrt(Dh)

        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(q.dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

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

    # ------------------------------------------------------------------
    # Compressive memory
    # ------------------------------------------------------------------
    def _compressive_memory_read(self, q, k, v, write_gate, valid_mask=None):
        """
        Sequential O(L) memory update + read.

        M_t = M_{t-1} + g_t · φ(k_t)^T v_t
        z_t = z_{t-1} + g_t · φ(k_t)
        y_t = φ(q_t) M_t / (φ(q_t) z_t + ε)

        INF3-FIX-04: the gate-scaled feature `g_t · φ(k_t)` is precomputed
        outside the loop, and each step uses a batched outer product.
        """
        B, H, L, Dh = q.shape
        device = q.device
        dtype = q.dtype

        # Feature map φ(x) = ELU(x) + 1
        phi_q = F.elu(q.float()) + 1.0  # (B, H, L, Dh)
        phi_k = F.elu(k.float()) + 1.0  # (B, H, L, Dh)
        v_f = v.float()  # (B, H, L, Dh)

        # Padding mask
        if valid_mask is not None:
            m = valid_mask[:, None, :, None].to(phi_k.dtype)  # (B, 1, L, 1)
            phi_k = phi_k * m
            v_f = v_f * m

        # Precompute gate-scaled keys
        # write_gate: (B, H, L, 1) → kg: (B, H, L, Dh)
        kg = write_gate.to(phi_k.dtype) * phi_k  # (B, H, L, Dh)

        # Transient memory (fresh per forward)
        M = torch.zeros(B, H, Dh, Dh, device=device, dtype=torch.float32)
        Z = torch.zeros(B, H, Dh, device=device, dtype=torch.float32)

        outputs = []
        for t in range(L):
            k_t = kg[:, :, t]  # (B, H, Dh)
            v_t = v_f[:, :, t]  # (B, H, Dh)
            q_t = phi_q[:, :, t]  # (B, H, Dh)

            # Batched outer product + accumulation
            M = M + k_t.unsqueeze(-1) * v_t.unsqueeze(-2)  # (B, H, Dh, Dh)
            Z = Z + k_t  # (B, H, Dh)

            # Read
            o_t = torch.einsum("bhd,bhde->bhe", q_t, M)  # (B, H, Dh)
            denom = torch.einsum("bhd,bhd->bh", q_t, Z).clamp(min=1e-6)
            o_t = o_t / denom.unsqueeze(-1)
            outputs.append(o_t)

        out = torch.stack(outputs, dim=2)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
        return out.to(dtype)

    # ------------------------------------------------------------------
    # Forward — signature matches BERT/RoBERTa/DistilBERT
    # ------------------------------------------------------------------
    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]:

        # INF3-FIX-02: layer_idx from kwargs (some HF models pass it)
        kwargs.get("layer_idx", 0)

        # INF3-FIX-01: position_embeddings comes from kwargs, never positional
        position_embeddings = kwargs.get("position_embeddings")

        if encoder_hidden_states is not None:
            raise NotImplementedError("InfiniAttention does not support cross-attention.")

        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, L, _ = hidden_states.shape
        H, Dh = self.num_heads, self.head_dim
        device = hidden_states.device

        # ---- projections ----
        q = self.q_proj(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
        k = (
            self.k_proj(hidden_states)
            .view(B, L, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, L, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- RoPE ----
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

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

        # ---- valid mask ----
        valid_mask = self._prepare_valid_mask(attention_mask, B, L, device)

        # ---- attention path ----
        if self.is_decoder:
            A_local, attn = self._local_sliding_attention(q, k, v, valid_mask)

            # Write gate: (B, L, H) → (B, H, L, 1)
            wg = torch.sigmoid(self.w_mem(hidden_states))
            write_gate = wg.transpose(1, 2).unsqueeze(-1)

            A_mem = self._compressive_memory_read(q, k, v, write_gate, valid_mask)

            beta = torch.sigmoid(self.beta).to(A_mem.dtype).view(1, H, 1, 1)
            fused = beta * A_mem + (1.0 - beta) * A_local
        else:
            fused, attn = self._full_attention(q, k, v, valid_mask)

        # ---- head_mask ----
        if head_mask is not None:
            hm = head_mask.to(fused.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, H, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], H, 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            # else assume 4D already
            fused = fused * hm

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

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

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

        # Zero out padded query positions
        if valid_mask is not None:
            out = out * valid_mask.unsqueeze(-1).to(out.dtype)

        return out, (attn if output_attentions else None)

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("q_proj", self.q_proj),
            ("k_proj", self.k_proj),
            ("v_proj", self.v_proj),
            ("out_proj", self.out_proj),
            ("w_mem", self.w_mem),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["beta"] = _sz(self.beta)
        out["total"] = sum(out.values())
        return out

Infini-Attention: sliding-window local attention + compressive memory.

For decoders, the output is a per-head convex combination of: - A_local: sliding-window softmax attention - A_mem: read from a transient (per-forward) compressive memory state For encoders, it degenerates to full softmax attention (memory is skipped).

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.

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 storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("q_proj", self.q_proj),
        ("k_proj", self.k_proj),
        ("v_proj", self.v_proj),
        ("out_proj", self.out_proj),
        ("w_mem", self.w_mem),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["beta"] = _sz(self.beta)
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class InfoMambaConceptBottleneck (dim: int, bottleneck_dim: int = 256)

Expand source code

class InfoMambaConceptBottleneck(nn.Module):
    def __init__(self, dim: int, bottleneck_dim: int = 256) -> None:
        super().__init__()
        self.dim = dim
        self.bottleneck_dim = bottleneck_dim
        self.norm_in = nn.LayerNorm(dim)
        self.W_down = nn.Linear(dim, bottleneck_dim, bias=False)
        self.norm_mid = nn.LayerNorm(bottleneck_dim)
        self.W_up = nn.Linear(bottleneck_dim, dim, bias=False)
        self.norm_out = nn.LayerNorm(dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.norm_in(x)
        z = self.W_down(h)
        z = self.norm_mid(z)
        z = F.gelu(z)
        h = self.W_up(z)
        return self.norm_out(h)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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:
    h = self.norm_in(x)
    z = self.W_down(h)
    z = self.norm_mid(z)
    z = F.gelu(z)
    h = self.W_up(z)
    return self.norm_out(h)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class Jamba2Block (config: Any | None = None,
d_model: int | None = None,
n_heads: int | None = None,
n_experts: int = 4,
top_k: int = 2,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "jamba", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "jamba2", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
class Jamba2Block(BaseAttention):
    r"""Jamba-2 Hybrid: SVD-Mamba + Multi-Head Attention + Expert-Choice MoE."""

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        d_model: int | None = None,
        n_heads: int | None = None,
        n_experts: int = 4,
        top_k: int = 2,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        self.config = config

        # ---------- d_model / n_heads ----------
        d_model = kwargs.get("d_model") or kwargs.get("hidden_size") or kwargs.get("dim") or d_model
        n_heads = (
            kwargs.get("n_heads")
            or kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or n_heads
        )

        if d_model is None and config is not None:
            d_model = getattr(
                config,
                "hidden_size",
                getattr(config, "d_model", getattr(config, "dim", 768)),
            )
        if d_model is None and source_attn is not None:
            for attr in ("out_lin", "dense", "c_proj", "o_proj", "out_proj"):
                layer = getattr(source_attn, attr, None)
                if layer is not None and hasattr(layer, "out_features"):
                    d_model = layer.out_features
                    break
        if d_model is None:
            d_model = 768

        if n_heads is None and config is not None:
            n_heads = getattr(config, "num_attention_heads", getattr(config, "n_heads", 12))
        if n_heads is None and source_attn is not None:
            n_heads = getattr(
                source_attn, "num_attention_heads", getattr(source_attn, "n_heads", 12)
            )
        if n_heads is None:
            n_heads = 12

        if config is not None:
            n_experts = int(getattr(config, "num_experts", getattr(config, "n_experts", n_experts)))
            top_k = int(getattr(config, "top_k", top_k))

        self.d_model = int(d_model)
        self.dim = self.d_model
        self.hidden_size = self.d_model
        self.n_heads = int(n_heads)
        self.num_heads = self.n_heads

        num_kv_heads = None
        if config is not None:
            num_kv_heads = getattr(
                config, "num_key_value_heads", getattr(config, "num_kv_heads", None)
            )
        if num_kv_heads is None:
            num_kv_heads = kwargs.get(
                "num_kv_heads", kwargs.get("num_key_value_heads", self.num_heads)
            )
        self.num_kv_heads = max(1, min(int(num_kv_heads), self.num_heads))

        self.head_dim = self.d_model // self.num_heads if self.num_heads > 0 else 64
        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 // self.num_kv_heads)
        self.supports_gqa = True

        self.is_decoder = _detect_decoder(config, kwargs.get("is_decoder"), source_attn, kwargs)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False
        self._return_format: str = "llama"  # updated after copy_weights

        # ---------- LoRA config ----------
        if config is not None:
            self.lora_rank = int(
                getattr(
                    config,
                    "lora_rank",
                    getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
                )
            )
            self.lora_alpha = float(
                getattr(
                    config,
                    "lora_alpha",
                    getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
                )
            )
        elif comp_config is not None:
            self.lora_rank = int(getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)))
            self.lora_alpha = float(
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
            )
        else:
            self.lora_rank = int(kwargs.get("lora_rank", 16))
            self.lora_alpha = float(kwargs.get("lora_alpha", 32.0))
        self.freeze_base = self.lora_rank > 0

        self.source_attn = source_attn
        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )

        # ---------- MHA projections ----------
        self.q_lin = LoRALinear(
            self.d_model,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_lin = LoRALinear(
            self.d_model,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_lin = LoRALinear(
            self.d_model,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_lin = LoRALinear(
            self.all_head_size,
            self.d_model,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.q_proj = self.q_lin
        self.k_proj = self.k_lin
        self.v_proj = self.v_lin
        self.out_proj = self.out_lin

        # ---------- SSM branch ----------
        self.mamba = SVDMambaSub(self.d_model, config=config, rank=rank, is_decoder=self.is_decoder)

        # ---------- MoE branch ----------
        self.moe = ExpertChoiceMoE(self.d_model, n_experts=n_experts, capacity_factor=1.0)

        # ---------- Norms ----------
        self.norm1 = nn.LayerNorm(self.d_model)
        self.norm2 = nn.LayerNorm(self.d_model)
        self.norm3 = nn.LayerNorm(self.d_model)

        dropout_p = float(
            getattr(
                config,
                "dropout",
                getattr(config, "attention_probs_dropout_prob", 0.1) if config else 0.1,
            )
        )
        self.dropout = nn.Dropout(dropout_p)

        # JAMBA-FIX-21: aux_loss accumulation buffer
        self._aux_loss_sum: float = 0.0
        self._aux_loss_count: int = 0

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

        # JAMBA-FIX-19: finalize return format
        self._return_format = _detect_return_format(source_attn, self.is_nested_encoder)

    # ------------------------------------------------------------------
    @property
    def aux_loss(self) -> float:
        """Mean aux loss over the accumulated forward calls."""
        if self._aux_loss_count == 0:
            return 0.0
        return self._aux_loss_sum / self._aux_loss_count

    def reset_aux_loss(self) -> None:
        """Call at the end of an epoch / accumulation window."""
        self._aux_loss_sum = 0.0
        self._aux_loss_count = 0

    # ------------------------------------------------------------------
    def freeze_base_weights(self) -> None:
        for mod in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
            if hasattr(mod, "freeze_base_weights"):
                mod.freeze_base_weights()

    def unfreeze_base_weights(self) -> None:
        for mod in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
            if hasattr(mod, "unfreeze_base_weights"):
                mod.unfreeze_base_weights()

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

        cls_name = getattr(src, "__class__", type(src)).__name__
        has_out = any(hasattr(src, a) for a in ("dense", "out_proj", "out_lin", "c_proj", "o_proj"))
        if cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(src, "query") and not has_out
        ):
            self.is_nested_encoder = True
        if hasattr(src, "self") and hasattr(src, "output"):
            self.is_nested_encoder = True

        extracted = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c = src.c_attn
            base_c = getattr(c, "base_layer", c)
            cw = getattr(base_c, "weight", getattr(c, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                is_conv = (
                    c.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or hasattr(c, "nf")
                    or (cw.ndim == 2 and cw.shape[0] < cw.shape[1])
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )

            if hasattr(src, "c_proj") and src.c_proj is not None:
                o = src.c_proj
                base_o = getattr(o, "base_layer", o)
                ow = getattr(base_o, "weight", getattr(o, "weight", None))
                if ow is not None:
                    ow = ow.data.clone()
                    is_conv_o = (
                        o.__class__.__name__ == "Conv1D"
                        or base_o.__class__.__name__ == "Conv1D"
                        or hasattr(o, "nf")
                        or (ow.ndim == 2 and ow.shape[0] < ow.shape[1])
                    )
                    if is_conv_o:
                        ow = ow.t().contiguous()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # Nested (BERT / RoBERTa)
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                d = src.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        # Flat (DistilBERT, Qwen, LLaMA)
        else:
            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 key, names in candidates.items():
                for n in names:
                    if hasattr(src, n):
                        m = getattr(src, n)
                        if m is not None and hasattr(m, "weight"):
                            w_data = m.weight.data.clone()
                            is_conv = (
                                m.__class__.__name__ == "Conv1D"
                                or hasattr(m, "nf")
                                or (
                                    key == "o"
                                    and w_data.ndim == 2
                                    and w_data.shape[0] < w_data.shape[1]
                                )
                            )
                            if is_conv:
                                w_data = w_data.t().contiguous()
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                w_data,
                                b.data.clone() if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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("Jamba2 UniversalAdapter fallback: %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"Jamba2Block copy_weights failed: missing mandatory "
                f"projections {missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_lin,
            "k": self.k_lin,
            "v": self.v_lin,
            "o": self.out_lin,
        }

        with torch.no_grad():
            for key, target in target_map.items():
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target.weight)
                        if target.bias is not None:
                            nn.init.zeros_(target.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target.weight.shape
                    and w.t().contiguous().shape == target.weight.shape
                ):
                    w = w.t().contiguous()

                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))
                target.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                )

                if b is not None and target.bias is not None:
                    mb = min(target.bias.numel(), b.numel())
                    target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))
                elif target.bias is not None:
                    nn.init.zeros_(target.bias)

                n = torch.linalg.norm(target.weight.data.float()).item()
                if n < 1e-4:
                    raise RuntimeError(f"Jamba2: weight for '{key}' norm too small ({n:.2e})")

                if self.freeze_base:
                    target.freeze_base_weights()

        if hasattr(src, "rotary_emb") and src.rotary_emb is not None:
            self.rotary_emb = src.rotary_emb

        if self.freeze_base:
            self.freeze_base_weights()

    # ------------------------------------------------------------------
    def _norm_dtype(self, norm: nn.LayerNorm, fallback: torch.dtype) -> torch.dtype:
        return norm.weight.dtype if norm.weight is not None else fallback

    def _format_output(
        self,
        out: torch.Tensor,
        attn_weights: torch.Tensor | None,
        present: Any | None,
    ) -> tuple:
        """JAMBA-FIX-19: emit the tuple layout expected by the host model."""
        fmt = self._return_format
        if fmt == "bert":
            return (out, attn_weights) if attn_weights is not None else (out,)
        if fmt == "gpt2":
            if attn_weights is not None:
                return (out, present, attn_weights)
            return (out, present)
        # default 'llama'
        if attn_weights is not None:
            return (out, attn_weights, present)
        return (out, present)

    def _normalize_head_mask(self, head_mask: torch.Tensor, attn: torch.Tensor) -> torch.Tensor:
        """JAMBA-FIX-22: normalise head_mask to (B, H, Lq, Lk) broadcast."""
        hm = head_mask.to(attn.dtype)
        if hm.dim() == 1:
            hm = hm.view(1, self.num_heads, 1, 1)
        elif hm.dim() == 2:
            if hm.shape[-1] == self.num_heads:
                hm = hm.view(hm.shape[0], self.num_heads, 1, 1)
            else:
                # (num_layers, num_heads) — take first row as a safe fallback
                hm = hm.view(1, hm.shape[-1], 1, 1)
        elif hm.dim() == 3:
            # (B, H, Lk) → broadcast over query positions
            hm = hm.unsqueeze(2)  # (B, H, 1, Lk)
        # dim == 4 → pass through
        return attn * hm

    # ------------------------------------------------------------------
    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,
    ):
        B, L, _ = hidden_states.shape
        layer_idx = int(kwargs.get("layer_idx", 0))
        position_embeddings = kwargs.get("position_embeddings")
        target_dtype = hidden_states.dtype

        # ---- Read cache ----
        past_k, past_v, past_len, past_mamba = _read_cache(past_key_value, layer_idx)

        # ---- Branch 1: SVD-Mamba (JAMBA-FIX-20) ----
        mamba_out, present_mamba = self.mamba(
            hidden_states,
            attention_mask=attention_mask,
            past_state=past_mamba,
        )
        n1d = self._norm_dtype(self.norm1, target_dtype)
        x1 = self.norm1((hidden_states + mamba_out.to(target_dtype)).to(n1d)).to(target_dtype)

        # ---- Branch 2: MHA ----
        q = self.q_lin(x1).view(B, L, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k = self.k_lin(x1).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        v = self.v_lin(x1).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()

        # ---- RoPE with offset ----
        if position_embeddings is not None:
            cos, sin = position_embeddings
            # Slice the position range [past_len, past_len + L)
            cos = cos.narrow(-2, past_len, L) if cos.shape[-2] >= past_len + L else cos
            sin = sin.narrow(-2, past_len, L) if sin.shape[-2] >= past_len + L else sin
            q, k = apply_rotary_pos_emb(q, k, cos, sin)
        elif self.rotary_emb is not None:
            pos_ids = kwargs.get("position_ids")
            if pos_ids is None:
                pos_ids = (
                    torch.arange(past_len, past_len + L, device=hidden_states.device)
                    .unsqueeze(0)
                    .expand(B, -1)
                )
            cos = sin = None
            for call in (
                lambda: self.rotary_emb(k, seq_len=L),
                lambda: self.rotary_emb(k, pos_ids),
                lambda: self.rotary_emb(pos_ids),
            ):
                try:
                    cos, sin = call()
                    break
                except Exception:
                    continue
            if cos is not None and sin is not None:
                q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=pos_ids)

        # ---- Concatenate cache (before repeat_kv to keep it compact) ----
        if past_k is not None:
            k = torch.cat([past_k.to(k.dtype), k], dim=-2)
            v = torch.cat([past_v.to(v.dtype), v], dim=-2)

        present_key_value = _write_cache(k, v, present_mamba, past_key_value, layer_idx)

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

        scale = 1.0 / math.sqrt(self.head_dim)
        is_causal = bool(self.is_decoder and L > 1)

        # ---- SDPA fast path (prefill only) ----
        use_sdpa = (
            head_mask is None and not output_attentions and attention_mask is None and past_len == 0
        )
        if use_sdpa:
            try:
                _qs = q.half() if q.dtype == torch.float32 else q
                _ks = k_exp.half() if k_exp.dtype == torch.float32 else k_exp
                _vs = v_exp.half() if v_exp.dtype == torch.float32 else v_exp
                _sdpa_out = F.scaled_dot_product_attention(
                    _qs,
                    _ks,
                    _vs,
                    is_causal=is_causal,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                attn_out = (
                    _sdpa_out.to(target_dtype)
                    .transpose(1, 2)
                    .contiguous()
                    .view(B, L, self.all_head_size)
                )
                attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)

                attn_proj = attn_out if self.is_nested_encoder else self.out_lin(attn_out)
                n2d = self._norm_dtype(self.norm2, target_dtype)
                x2 = self.norm2((x1 + self.dropout(attn_proj.to(target_dtype))).to(n2d)).to(
                    target_dtype
                )

                moe_out = self.moe(x2)
                if isinstance(moe_out, (tuple, list)):
                    moe_out = moe_out[0]
                n3d = self._norm_dtype(self.norm3, target_dtype)
                x3 = self.norm3((x2 + self.dropout(moe_out.to(target_dtype))).to(n3d)).to(
                    target_dtype
                )

                self._accumulate_aux_loss()
                return self._format_output(x3, None, present_key_value)
            except Exception as e:
                logger.debug("Jamba2 SDPA path failed: %s", e)

        # ---- Manual attention ----
        scores = torch.matmul(q.float(), k_exp.float().transpose(-2, -1).contiguous()) * scale

        if is_causal:
            kv_len = k_exp.shape[-2]
            causal = torch.tril(
                torch.ones(L, kv_len, device=hidden_states.device, dtype=torch.bool),
                diagonal=past_len,
            )
            scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

        if attention_mask is not None:
            if attention_mask.dim() == 2:
                valid = get_token_valid_mask(
                    attention_mask, B, k_exp.shape[-2], hidden_states.device
                )
                if valid is not None:
                    scores = scores.masked_fill(~valid.unsqueeze(1).unsqueeze(2), -1e4)
            elif attention_mask.dim() == 3:
                scores = scores + attention_mask.unsqueeze(1).float()
            elif attention_mask.dim() == 4:
                if attention_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~attention_mask, -1e4)
                else:
                    scores = scores + attention_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(v_exp.dtype)
        probs = torch.nan_to_num(probs, nan=0.0)
        probs = self.dropout(probs)

        # JAMBA-FIX-22: normalized head_mask
        if head_mask is not None:
            probs = self._normalize_head_mask(head_mask, probs)

        attn_out = (
            torch.matmul(probs, v_exp).transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        )
        attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)

        attn_proj = attn_out if self.is_nested_encoder else self.out_lin(attn_out)
        n2d = self._norm_dtype(self.norm2, target_dtype)
        x2 = self.norm2((x1 + self.dropout(attn_proj.to(target_dtype))).to(n2d)).to(target_dtype)

        moe_out = self.moe(x2)
        if isinstance(moe_out, (tuple, list)):
            moe_out = moe_out[0]
        n3d = self._norm_dtype(self.norm3, target_dtype)
        x3 = self.norm3((x2 + self.dropout(moe_out.to(target_dtype))).to(n3d)).to(target_dtype)

        self._accumulate_aux_loss()
        attn_weights = probs if output_attentions else None
        return self._format_output(x3, attn_weights, present_key_value)

    def _accumulate_aux_loss(self) -> None:
        """JAMBA-FIX-21: rolling mean of the MoE auxiliary loss."""
        aux = getattr(self.moe, "aux_loss", 0.0)
        v = float(aux.item()) if torch.is_tensor(aux) else float(aux)
        self._aux_loss_sum += v
        self._aux_loss_count += 1

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p: Any) -> int:
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out: dict[str, int] = {}
        for name, m in (
            ("q_lin", self.q_lin),
            ("k_lin", self.k_lin),
            ("v_lin", self.v_lin),
            ("out_lin", self.out_lin),
            ("mamba", self.mamba),
            ("moe", self.moe),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["total"] = sum(out.values())
        return out

Jamba-2 Hybrid: SVD-Mamba + Multi-Head Attention + Expert-Choice MoE.

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.

Instance variables

prop aux_loss : float

Expand source code

@property
def aux_loss(self) -> float:
    """Mean aux loss over the accumulated forward calls."""
    if self._aux_loss_count == 0:
        return 0.0
    return self._aux_loss_sum / self._aux_loss_count

Mean aux loss over the accumulated forward calls.

Methods

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    for mod in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
        if hasattr(mod, "freeze_base_weights"):
            mod.freeze_base_weights()

def reset_aux_loss(self) ‑> None

Expand source code

def reset_aux_loss(self) -> None:
    """Call at the end of an epoch / accumulation window."""
    self._aux_loss_sum = 0.0
    self._aux_loss_count = 0

Call at the end of an epoch / accumulation window.

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p: Any) -> int:
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out: dict[str, int] = {}
    for name, m in (
        ("q_lin", self.q_lin),
        ("k_lin", self.k_lin),
        ("v_lin", self.v_lin),
        ("out_lin", self.out_lin),
        ("mamba", self.mamba),
        ("moe", self.moe),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["total"] = sum(out.values())
    return out

def unfreeze_base_weights(self) ‑> None

Expand source code

def unfreeze_base_weights(self) -> None:
    for mod in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
        if hasattr(mod, "unfreeze_base_weights"):
            mod.unfreeze_base_weights()

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class JambaInfoBlock (config: Any | None = None,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
is_decoder: bool | None = None,
per_head_alpha: bool = True,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("jamba_info", supports_causal=True, supports_gqa=True)
class JambaInfoBlock(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        is_decoder: bool | None = None,
        per_head_alpha: bool = True,
        **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.is_decoder = _detect_decoder(config, is_decoder, source_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False
        self._return_format = "llama"

        # LoRA
        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.freeze_base = self.lora_rank > 0

        # ---------- Projections ----------
        def _mk_qk():
            return LoRALinear(
                self.dim,
                self.all_head_size,
                bias=False,
                lora_rank=self.lora_rank,
                lora_alpha=self.lora_alpha,
                freeze_base=self.freeze_base,
            )

        def _mk_v():
            return LoRALinear(
                self.dim,
                self.all_kv_head_size,
                bias=False,
                lora_rank=self.lora_rank,
                lora_alpha=self.lora_alpha,
                freeze_base=self.freeze_base,
            )

        self.W_q1 = _mk_qk()
        self.W_k1 = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_q2 = _mk_qk()
        self.W_k2 = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.W_v = _mk_v()
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Aliases
        self.q_proj = self.W_q1
        self.k_proj = self.W_k1
        self.v_proj = self.W_v

        # ---------- JI-02: learnable per-head λ (DIFF Transformer style) ----------
        # λ_full[h] = exp(λ_q1[h] · λ_k1[h]) - exp(λ_q2[h] · λ_k2[h]) + λ_init[h]
        self.lambda_q1 = nn.Parameter(torch.zeros(self.num_heads))
        self.lambda_k1 = nn.Parameter(torch.zeros(self.num_heads))
        self.lambda_q2 = nn.Parameter(torch.zeros(self.num_heads))
        self.lambda_k2 = nn.Parameter(torch.zeros(self.num_heads))
        # λ_init per head; initial value mimics the previous constant 0.2
        self.lambda_init = nn.Parameter(torch.full((self.num_heads,), 0.2))

        # ---------- JI-07: per-head α ----------
        self.per_head_alpha = bool(per_head_alpha)
        if self.per_head_alpha:
            self.alpha_raw = nn.Parameter(torch.full((self.num_heads,), -2.0))
        else:
            self.alpha_raw = nn.Parameter(torch.tensor(-2.0))

        # ---------- Bottleneck + SSM ----------
        bottleneck_dim = int(getattr(comp_config, "bottleneck_dim", 256) if comp_config else 256)
        self.bottleneck = InfoMambaConceptBottleneck(self.dim, bottleneck_dim=bottleneck_dim)

        self.ssm_fwd: nn.Module | None = None
        self.ssm_bwd: nn.Module | None = None  # only for encoders (JI-10)
        if SimpleSSM is not None:
            self.ssm_fwd = self._make_ssm(config)
            if not self.is_decoder:
                self.ssm_bwd = self._make_ssm(config)

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

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

        self._return_format = _detect_return_format(source_attn, self.is_nested_encoder)

    # ------------------------------------------------------------------
    def _make_ssm(self, config: Any | None) -> nn.Module | None:
        if SimpleSSM is None:
            return None
        try:
            return SimpleSSM(d_model=self.dim, config=config, is_decoder=self.is_decoder)
        except TypeError:
            try:
                return SimpleSSM(d_model=self.dim, config=config)
            except TypeError:
                try:
                    return SimpleSSM(self.dim)
                except Exception as e:
                    logger.warning("SimpleSSM instantiation failed: %s", e)
                    return None

    @staticmethod
    def _detect_hidden_size(config, src, kwargs) -> 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", "W_q1", "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, src, kwargs, hidden_size) -> 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, src, kwargs, num_heads) -> 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, src, kwargs, hidden_size, num_heads) -> 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", "query", "q_lin", "W_q1"):
                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))

    # ---------------- Misc ----------------
    def _run_ssm(
        self,
        ssm: nn.Module | None,
        x: torch.Tensor,
        attention_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        if ssm is None:
            return torch.zeros_like(x)
        try:
            out = ssm(x, attention_mask=attention_mask)
        except TypeError:
            try:
                out = ssm(x)
            except Exception as e:
                logger.debug("SSM forward failed: %s", e)
                return torch.zeros_like(x)
        if isinstance(out, tuple):
            out = out[0]
        return out

    def _format_output(
        self, out: torch.Tensor, attn_weights: torch.Tensor | None, present: Any | None
    ) -> tuple:
        fmt = self._return_format
        if fmt == "bert":
            return (out, attn_weights)
        if fmt == "gpt2":
            return (out, present, attn_weights)
        return (out, attn_weights, present)

    def _normalize_head_mask(self, head_mask: torch.Tensor, attn: torch.Tensor) -> torch.Tensor:
        """JI-03: robust normalisation. attn: (B, H, Lq, Lk)."""
        hm = head_mask.to(attn.dtype)
        H = self.num_heads
        if hm.dim() == 1:
            if hm.numel() == H:
                hm = hm.view(1, H, 1, 1)
            else:
                hm = hm.view(1, -1, 1, 1)
        elif hm.dim() == 2:
            # Prefer interpreting last dim as heads; else treat as (layers, heads)
            if hm.shape[-1] == H:
                hm = hm.view(hm.shape[0], H, 1, 1)
            elif hm.shape[0] == H:
                hm = hm.t().view(1, H, 1, 1)
            else:
                hm = hm.view(1, hm.shape[-1], 1, 1)
        elif hm.dim() == 3:
            # (B, H, Lk) or (B, Lk, H)
            if hm.shape[1] == H:
                hm = hm.unsqueeze(2)  # (B, H, 1, Lk)
            elif hm.shape[2] == H:
                hm = hm.permute(0, 2, 1).unsqueeze(2)  # (B, H, 1, Lk)
            else:
                hm = hm.unsqueeze(1)
        # dim == 4: pass through
        return attn * hm

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

        cls_name = src.__class__.__name__
        has_out_lin = any(hasattr(src, a) for a in ("out_lin", "c_proj", "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 = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c = src.c_attn
            base_c = getattr(c, "base_layer", c)
            cw = getattr(base_c, "weight", getattr(c, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                is_conv = (
                    c.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (cw.ndim == 2 and cw.shape[1] == 3 * self.hidden_size)
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )
            if hasattr(src, "c_proj") and src.c_proj is not None:
                o = src.c_proj
                base_o = getattr(o, "base_layer", o)
                ow = getattr(base_o, "weight", getattr(o, "weight", None))
                if ow is not None:
                    ow = ow.data.clone()
                    is_conv_o = (
                        o.__class__.__name__ == "Conv1D"
                        or base_o.__class__.__name__ == "Conv1D"
                        or (ow.ndim == 2 and ow.shape[0] < ow.shape[1])
                    )
                    if is_conv_o:
                        ow = ow.t().contiguous()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                d = src.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data.clone(),
                    b.data.clone() if b is not None else None,
                )

        else:
            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 key, names in candidates.items():
                for n in names:
                    if hasattr(src, n):
                        m = getattr(src, n)
                        if m is not None and hasattr(m, "weight"):
                            w = m.weight.data.clone()
                            is_conv = m.__class__.__name__ == "Conv1D" or (
                                key == "o" and w.ndim == 2 and w.shape[0] < w.shape[1]
                            )
                            if is_conv:
                                w = w.t().contiguous()
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                w,
                                b.data.clone() if b 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: %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:
            logger.warning(
                "JambaInfoBlock.copy_weights: missing %s from %s; keeping defaults.",
                missing_after,
                src.__class__.__name__,
            )

        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

        def _copy_into(
            target_mod: LoRALinear, w: torch.Tensor | None, b: torch.Tensor | None
        ) -> None:
            if w is None:
                return
            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)
            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            r = min(target_mod.weight.size(0), w.size(0))
            c = min(target_mod.weight.size(1), w.size(1))
            target_mod.weight.data[:r, :c].copy_(
                w[:r, :c].to(device=target_mod.weight.device, dtype=target_mod.weight.dtype)
            )
            if b is not None and target_mod.bias is not None:
                n = min(b.numel(), target_mod.bias.numel())
                target_mod.bias.data[:n].copy_(
                    b[:n].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )

        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:
                _copy_into(self.W_q1, _orient(w_q, self.hidden_size), None)
                _copy_into(self.W_q2, _orient(w_q, self.hidden_size), None)
            if w_k is not None:
                _copy_into(self.W_k1, _orient(w_k, self.hidden_size), None)
                _copy_into(self.W_k2, _orient(w_k, self.hidden_size), None)
            if w_v is not None:
                _copy_into(self.W_v, _orient(w_v, self.hidden_size), None)
            if w_o is not None:
                _copy_into(self.out_proj, _orient(w_o, self.all_head_size), b_o)
            elif self.is_nested_encoder:
                nn.init.eye_(self.out_proj.weight)
                if self.out_proj.bias is not None:
                    nn.init.zeros_(self.out_proj.bias)

        # JI-11: only warn, do not reinit. Reinit destroys legit small weights.
        for name, mod in (
            ("W_q1", self.W_q1),
            ("W_k1", self.W_k1),
            ("W_q2", self.W_q2),
            ("W_k2", self.W_k2),
            ("W_v", self.W_v),
            ("out_proj", self.out_proj),
        ):
            w_norm = float(torch.norm(mod.weight.data).item())
            if w_norm < 1e-6:
                logger.warning(
                    "JambaInfoBlock: weight '%s' has near-zero norm (%.2e). "
                    "Verify the source attention was properly initialized.",
                    name,
                    w_norm,
                )

        if self.freeze_base:
            for mod in (
                self.W_q1,
                self.W_k1,
                self.W_q2,
                self.W_k2,
                self.W_v,
                self.out_proj,
            ):
                mod.freeze_base_weights()

    # ---------------- 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,
        position_ids: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        *args: Any,
        **kwargs: Any,
    ):
        # Mask may come via kwargs
        mask = kwargs.get("mask")
        if attention_mask is None and mask is not None:
            attention_mask = mask

        x = hidden_states
        B, L, D = x.shape
        device = x.device

        # --- JI-01: read cache ---
        past_k1, past_k2, past_v, past_len = _read_cache(past_key_value)

        # --- QKV projections ---
        q1 = self.W_q1(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k1 = self.W_k1(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        q2 = self.W_q2(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k2 = self.W_k2(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        v = self.W_v(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()

        # --- JI-05: position_ids accepted; JI-06: RoPE offset ---
        position_embeddings = position_embeddings or kwargs.get("position_embeddings")
        if position_embeddings is None and position_ids is None:
            position_ids = kwargs.get("position_ids")

        if position_embeddings is not None:
            cos, sin = position_embeddings
            # Slice to current positions
            if cos.shape[-2] >= past_len + L:
                cos = cos.narrow(-2, past_len, L)
                sin = sin.narrow(-2, past_len, L)
            q1, k1 = apply_rotary_pos_emb(q1, k1, cos, sin)
            q2, k2 = apply_rotary_pos_emb(q2, k2, cos, sin)
        elif position_ids is not None:
            # RoPE via position_ids not supported here without a rotary_emb;
            # caller should pass position_embeddings.
            pass

        # --- JI-01: append to cache (before GQA expansion) ---
        if past_k1 is not None:
            k1 = torch.cat([past_k1.to(k1.dtype), k1], dim=-2)
            k2 = torch.cat([past_k2.to(k2.dtype), k2], dim=-2)
            v = torch.cat([past_v.to(v.dtype), v], dim=-2)
        present_key_value = (k1, k2, v)

        # --- GQA expand ---
        if self.kv_repeat > 1:
            k1 = repeat_kv(k1, self.kv_repeat)
            k2 = repeat_kv(k2, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        scale = 1.0 / math.sqrt(self.head_dim)
        scores1 = torch.matmul(q1.float(), k1.float().transpose(-2, -1).contiguous()) * scale
        scores2 = torch.matmul(q2.float(), k2.float().transpose(-2, -1).contiguous()) * scale

        # --- Causal mask with past_len offset ---
        if self.is_decoder and L > 1:
            kv_len = k1.shape[-2]
            causal = torch.tril(
                torch.ones(L, kv_len, device=device, dtype=torch.bool),
                diagonal=past_len,
            )
            scores1 = scores1.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)
            scores2 = scores2.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

        # --- JI-04: attention_mask dims 2/3/4 ---
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                valid = get_token_valid_mask(attention_mask, B, k1.shape[-2], device)
                if valid is not None:
                    scores1 = scores1.masked_fill(~valid.unsqueeze(1).unsqueeze(2), -1e4)
                    scores2 = scores2.masked_fill(~valid.unsqueeze(1).unsqueeze(2), -1e4)
            elif attention_mask.dim() == 3:
                # (B, Lq, Lk) broadcast to (B, 1, Lq, Lk)
                m3 = attention_mask.unsqueeze(1).float()
                if attention_mask.dtype == torch.bool:
                    m3 = (~attention_mask).unsqueeze(1).float() * -1e4
                scores1 = scores1 + m3
                scores2 = scores2 + m3
            elif attention_mask.dim() == 4:
                if attention_mask.dtype == torch.bool:
                    m4 = (~attention_mask).float() * -1e4
                else:
                    m4 = attention_mask.float()
                scores1 = scores1 + m4
                scores2 = scores2 + m4

        scores1 = torch.nan_to_num(scores1, nan=-1e4, posinf=1e4, neginf=-1e4)
        scores2 = torch.nan_to_num(scores2, nan=-1e4, posinf=1e4, neginf=-1e4)

        attn1 = F.softmax(scores1, dim=-1).to(q1.dtype)
        attn2 = F.softmax(scores2, dim=-1).to(q2.dtype)

        # --- JI-02: learnable per-head λ ---
        # λ_full[h] = exp(λ_q1[h] · λ_k1[h]) - exp(λ_q2[h] · λ_k2[h]) + λ_init[h]
        lambda_full = (
            torch.exp(self.lambda_q1 * self.lambda_k1)
            - torch.exp(self.lambda_q2 * self.lambda_k2)
            + self.lambda_init
        ).to(attn1.dtype)  # (H,)
        lambda_full = lambda_full.view(1, self.num_heads, 1, 1)

        diff = attn1 - lambda_full * attn2
        diff = F.relu(diff)
        diff = diff / (diff.sum(dim=-1, keepdim=True) + 1e-8)

        diff_attn = self.dropout(diff).to(v.dtype)

        # --- JI-03: head_mask ---
        if head_mask is not None:
            diff_attn = self._normalize_head_mask(head_mask, diff_attn)

        attn_out = (
            torch.matmul(diff_attn, v).transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        )

        # --- Bottleneck + SSM ---
        bottleneck_x = self.bottleneck(x)

        if self.is_decoder:
            ssm_branch = self._run_ssm(self.ssm_fwd, bottleneck_x, attention_mask)
        else:
            # JI-10: use a separate backward SSM if available
            if self.ssm_bwd is not None:
                x_rev = torch.flip(bottleneck_x, dims=[1])
                ssm_fwd = self._run_ssm(self.ssm_fwd, bottleneck_x, attention_mask)
                ssm_bwd = self._run_ssm(self.ssm_bwd, x_rev, attention_mask)
                ssm_branch = 0.5 * (ssm_fwd + torch.flip(ssm_bwd, dims=[1]))
            else:
                ssm_branch = self._run_ssm(self.ssm_fwd, bottleneck_x, attention_mask)

        # --- JI-07: α combination (per-head or scalar) ---
        target_dtype = x.dtype
        if self.per_head_alpha:
            alpha = torch.sigmoid(self.alpha_raw).to(target_dtype)  # (H,)
            alpha = alpha.view(1, self.num_heads, 1, 1)
            attn_out_h = attn_out.view(B, L, self.num_heads, self.head_dim)
            attn_out_h = attn_out_h.transpose(1, 2)  # (B, H, L, Dh)
            ssm_branch_h = ssm_branch.view(B, L, self.num_heads, self.head_dim)
            ssm_branch_h = ssm_branch_h.transpose(1, 2)  # (B, H, L, Dh)
            combined = (1.0 - alpha) * attn_out_h.to(target_dtype) + alpha * ssm_branch_h.to(
                target_dtype
            )
            combined = combined.transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        else:
            alpha = torch.sigmoid(self.alpha_raw).to(target_dtype)
            combined = (1.0 - alpha) * attn_out.to(target_dtype) + alpha * ssm_branch.to(
                target_dtype
            )
        combined = torch.nan_to_num(combined, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            out = combined.to(target_dtype)
        else:
            out = self.out_proj(combined).to(target_dtype)

        attn_weights = diff_attn if output_attentions else None
        return self._format_output(out, attn_weights, present_key_value)

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class KDAAttention (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
num_heads: int | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("kda", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("kda_attention", supports_causal=True, supports_gqa=True)
class KDAAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        num_heads: int | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | 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 dim discovery ----
        if d_model is None:
            d_model = kwargs.get("d_model") or kwargs.get("hidden_size") or kwargs.get("dim")
            if d_model is None:
                if config is not None:
                    d_model = getattr(config, "hidden_size", getattr(config, "dim", 768))
                elif source_attn is not None:
                    out_layer = getattr(
                        source_attn,
                        "out_lin",
                        getattr(
                            source_attn,
                            "dense",
                            getattr(
                                source_attn,
                                "c_proj",
                                getattr(
                                    source_attn,
                                    "o_proj",
                                    getattr(source_attn, "out_proj", None),
                                ),
                            ),
                        ),
                    )
                    if out_layer is not None:
                        d_model = getattr(
                            out_layer,
                            "out_features",
                            getattr(out_layer, "in_features", 768),
                        )
                    else:
                        q_l = getattr(
                            source_attn,
                            "q_proj",
                            getattr(
                                source_attn,
                                "query",
                                getattr(source_attn, "q_lin", None),
                            ),
                        )
                        d_model = (
                            getattr(q_l, "in_features", getattr(q_l, "out_features", 768))
                            if q_l is not None
                            else 768
                        )
                else:
                    d_model = 768
        if hasattr(d_model, "hidden_size"):
            d_model = d_model.hidden_size

        # ---- num heads ----
        if num_heads is None:
            num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")
            if num_heads is None:
                if config is not None:
                    num_heads = getattr(
                        config, "num_attention_heads", getattr(config, "n_heads", 12)
                    )
                elif source_attn is not None:
                    num_heads = getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(source_attn, "n_heads", 12),
                    )
                else:
                    num_heads = 12

        self.dim = int(d_model)
        self.hidden_size = self.dim
        self.num_heads = int(num_heads)
        self.num_kv_heads = (
            getattr(
                config,
                "num_key_value_heads",
                getattr(config, "num_kv_heads", self.num_heads),
            )
            if config
            else self.num_heads
        )
        self.head_dim = self.dim // self.num_heads if self.num_heads > 0 else 64
        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

        # ---- decoder detection ----
        model_type = getattr(config, "model_type", "").lower() if config else ""
        is_dec = (
            getattr(
                config,
                "is_decoder",
                getattr(config, "is_causal", kwargs.get("is_causal", False)),
            )
            if config
            else False
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(_mt.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        # ---- KDA-FIX-02: rename to lora_* prefix for framework recognition ----
        # KDA-FIX-01: init at 5.0 → sigmoid(5) ≈ 0.9933 → effective context ~150 tokens
        self.lora_beta = nn.Parameter(torch.ones(self.num_heads, 1, 1) * 5.0)
        self.lora_output_scale = nn.Parameter(torch.ones(1))

        # ---- LoRA config ----
        if config is not None:
            self.lora_rank = getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
        elif comp_config is not None:
            self.lora_rank = getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16))
        else:
            self.lora_rank = kwargs.get("lora_rank", 16)

        if config is not None:
            self.lora_alpha = getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        elif comp_config is not None:
            self.lora_alpha = getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
        else:
            self.lora_alpha = kwargs.get("lora_alpha", 32.0)

        self.freeze_base = self.lora_rank > 0
        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )

        # ---- projections ----
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

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

    def freeze_base_weights(self) -> None:
        for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
            if hasattr(mod, "freeze_base_weights"):
                mod.freeze_base_weights()

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        if src is None:
            return

        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")
        )
        if cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(src, "query") and not has_out
        ):
            self.is_nested_encoder = True
        if hasattr(src, "self") and hasattr(src, "output"):
            self.is_nested_encoder = True

        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()
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or hasattr(c_attn_mod, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    c_w = c_w.t()
                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()
                    if (
                        c_proj_mod.__class__.__name__ == "Conv1D"
                        or base_o.__class__.__name__ == "Conv1D"
                        or hasattr(c_proj_mod, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        o_w = o_w.t()
                    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. 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
        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"):
                            w_data = mod.weight.data.clone()
                            if (
                                mod.__class__.__name__ == "Conv1D"
                                or hasattr(mod, "nf")
                                or (
                                    w_data.ndim == 2
                                    and proj_key == "o"
                                    and w_data.shape[0] < w_data.shape[1]
                                )
                            ):
                                w_data = w_data.t()
                            extracted[proj_key] = (
                                w_data,
                                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: %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"KDAAttention copy_weights failed: missing mandatory projections "
                f"{missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if w.shape != target_mod.weight.shape and target_mod.weight.shape == w.t().shape:
                    w = w.t()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                w_norm = float(torch.norm(w[:out_d, :in_d]).item())
                if w_norm < 1e-4:
                    raise RuntimeError(
                        f"KDAAttention weight verification failed for "
                        f"'{proj_key}': norm too small ({w_norm:.2e})"
                    )

                if self.freeze_base:
                    target_mod.freeze_base_weights()

        if self.lora_rank > 0 or self.freeze_base:
            self.freeze_base_weights()

    # ------------------------------------------------------------------
    # KDA-FIX-03: chunked scan with bidirectional guard
    # ------------------------------------------------------------------
    def _kda_scan_chunked(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        decay: torch.Tensor,
        mask: torch.Tensor | None = None,
        chunk_size: int = 256,
        _allow_bidirectional: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, H, L, d = q.shape
        gamma = decay.clamp(min=1e-4, max=1.0 - 1e-4)
        log_gamma = torch.log(gamma)

        mask_float = None
        if mask is not None:
            valid_m = get_token_valid_mask(mask, B, L, q.device)
            mask_float = valid_m.float()

        y_fwd_list = []
        S = torch.zeros(B, H, d, v.shape[-1], device=q.device, dtype=torch.float32)
        Z = torch.zeros(B, H, d, 1, device=q.device, dtype=torch.float32)

        num_chunks = (L + chunk_size - 1) // chunk_size
        for c in range(num_chunks):
            start = c * chunk_size
            end = min(start + chunk_size, L)
            qc = q[:, :, start:end, :].float()
            kc = k[:, :, start:end, :].float()
            vc = v[:, :, start:end, :].float()
            Lc = end - start

            mc = None
            if mask_float is not None:
                mc = mask_float[:, start:end]
                mc_col = mc.unsqueeze(1).unsqueeze(-1)
                kc = kc * mc_col
                vc = vc * mc_col

            pos = torch.arange(Lc, device=q.device, dtype=torch.float32)
            pos_diff = pos.unsqueeze(-1) - pos.unsqueeze(-2)
            tril_mask = torch.tril(torch.ones(Lc, Lc, device=q.device, dtype=torch.bool))
            D_intra = (
                torch.exp(
                    (pos_diff.unsqueeze(0).unsqueeze(0) * log_gamma).clamp(min=-40.0, max=0.0)
                )
                * tril_mask.float()
            )

            attn_intra = (qc @ kc.transpose(-1, -2)) * D_intra
            if mc is not None:
                mask_2d = mc.unsqueeze(1).unsqueeze(-1) * mc.unsqueeze(1).unsqueeze(-2)
                attn_intra = attn_intra * mask_2d

            num_intra = attn_intra @ vc
            den_intra = attn_intra.sum(dim=-1, keepdim=True)

            decay_inter = torch.exp(
                ((pos + 1).unsqueeze(0).unsqueeze(0).unsqueeze(-1) * log_gamma).clamp(
                    min=-40.0, max=0.0
                )
            )
            num_inter = (qc @ S) * decay_inter
            den_inter = (qc @ Z) * decay_inter

            if mc is not None:
                num_inter = num_inter * mc.unsqueeze(1).unsqueeze(-1)
                den_inter = den_inter * mc.unsqueeze(1).unsqueeze(-1)

            num_total = num_intra + num_inter
            den_total = (den_intra + den_inter).clamp(min=1e-4)
            yc = (num_total / den_total).to(q.dtype)
            if mc is not None:
                yc = yc * mc.unsqueeze(1).unsqueeze(-1)
            y_fwd_list.append(yc)

            decay_to_end = torch.exp(
                ((Lc - 1 - pos).unsqueeze(0).unsqueeze(0).unsqueeze(-1) * log_gamma).clamp(
                    min=-40.0, max=0.0
                )
            )
            k_decayed = kc * decay_to_end
            chunk_S = k_decayed.transpose(-1, -2) @ vc
            chunk_Z = k_decayed.transpose(-1, -2) @ torch.ones_like(vc[..., :1])

            gamma_chunk = torch.exp((Lc * log_gamma).clamp(min=-40.0, max=0.0))
            S = S * gamma_chunk + chunk_S
            Z = Z * gamma_chunk + chunk_Z

        y_fwd = torch.cat(y_fwd_list, dim=2)

        # KDA-FIX-03: guard against infinite recursion in bidirectional mode
        if (not self.is_decoder) and _allow_bidirectional:
            q_rev = q.flip(dims=[2])
            k_rev = k.flip(dims=[2])
            v_rev = v.flip(dims=[2])
            mask_rev = mask_float.flip(dims=[-1]) if mask_float is not None else None
            y_bwd_rev, _ = self._kda_scan_chunked(
                q_rev,
                k_rev,
                v_rev,
                decay,
                mask_rev,
                chunk_size=chunk_size,
                _allow_bidirectional=False,
            )
            y_bwd = y_bwd_rev.flip(dims=[2])
            return y_fwd, y_bwd

        return y_fwd, None

    # ------------------------------------------------------------------
    # Parallel scan (short sequences)
    # ------------------------------------------------------------------
    def _kda_scan_parallel(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        decay: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, H, L, d = q.shape

        if L > 512:
            return self._kda_scan_chunked(q, k, v, decay, mask, chunk_size=256)

        log_decay = torch.log(decay.clamp(min=1e-4, max=1.0 - 1e-4))
        pos = torch.arange(L, device=q.device, dtype=torch.float32)
        pos_diff = pos.unsqueeze(-1) - pos.unsqueeze(-2)

        tril_mask = torch.tril(torch.ones(L, L, device=q.device, dtype=torch.bool))
        D_causal = (
            torch.exp((pos_diff.unsqueeze(0).unsqueeze(0) * log_decay).clamp(min=-40.0, max=0.0))
            * tril_mask.float()
        )

        attn_scores = q.float() @ k.float().transpose(-1, -2)

        mask_2d = None
        if mask is not None:
            valid_m = get_token_valid_mask(mask, B, L, q.device)
            mask_float = valid_m.float()
            mask_2d = mask_float.unsqueeze(1).unsqueeze(-1) * mask_float.unsqueeze(1).unsqueeze(-2)
            D_causal = D_causal * mask_2d

        num_fwd = (attn_scores * D_causal) @ v.float()
        denom_fwd = (attn_scores * D_causal).sum(dim=-1, keepdim=True)
        y_fwd = torch.clamp(
            num_fwd / denom_fwd.clamp(min=1e-4),
            min=-1e4,
            max=1e4,
        ).to(q.dtype)

        if not self.is_decoder:
            triu_mask = torch.triu(torch.ones(L, L, device=q.device, dtype=torch.bool))
            D_anti = (
                torch.exp(
                    (-pos_diff.unsqueeze(0).unsqueeze(0) * log_decay).clamp(min=-40.0, max=0.0)
                )
                * triu_mask.float()
            )
            if mask_2d is not None:
                D_anti = D_anti * mask_2d
            num_bwd = (attn_scores * D_anti) @ v.float()
            denom_bwd = (attn_scores * D_anti).sum(dim=-1, keepdim=True)
            y_bwd = torch.clamp(
                num_bwd / denom_bwd.clamp(min=1e-4),
                min=-1e4,
                max=1e4,
            ).to(q.dtype)
            return y_fwd, y_bwd

        return y_fwd, None

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim

        q_raw = self.q_proj(hidden_states).view(B, L, H, d).transpose(1, 2).contiguous()
        k_raw = (
            self.k_proj(hidden_states).view(B, L, self.num_kv_heads, d).transpose(1, 2).contiguous()
        )
        v = self.v_proj(hidden_states).view(B, L, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # RoPE
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q_raw, k_raw = apply_rotary_pos_emb(q_raw, k_raw, cos, sin)
        elif self.rotary_emb is not None:
            rotary_module = self.rotary_emb
            pos_ids = kwargs.get("position_ids")
            if pos_ids is None:
                pos_ids = torch.arange(L, device=hidden_states.device).unsqueeze(0).expand(B, -1)
            try:
                cos, sin = rotary_module(q_raw, pos_ids)
                q_raw, k_raw = apply_rotary_pos_emb(q_raw, k_raw, cos, sin, position_ids=pos_ids)
            except Exception:
                try:
                    cos, sin = rotary_module(q_raw, seq_len=L)
                    q_raw, k_raw = apply_rotary_pos_emb(
                        q_raw, k_raw, cos, sin, position_ids=pos_ids
                    )
                except Exception as e:
                    logger.debug("KDA rotary_emb failed: %s", e)

        if self.kv_repeat > 1:
            k_raw = repeat_kv(k_raw, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        q = F.elu(q_raw) + 1.0
        k = (F.elu(k_raw) + 1.0) / (d**0.5)

        # KDA-FIX-01 + KDA-FIX-02: decay from lora_beta (trainable, init 5.0)
        decay = torch.sigmoid(self.lora_beta).view(1, H, 1, 1)

        mask_bool = None
        if attention_mask is not None:
            mask_bool = get_token_valid_mask(attention_mask, B, L, hidden_states.device)

        y_fwd, y_bwd = self._kda_scan_parallel(q, k, v, decay, mask_bool)

        if self.is_decoder or y_bwd is None:
            y = y_fwd
        else:
            y = 0.5 * (y_fwd + y_bwd)

        y = torch.nan_to_num(y, nan=0.0)

        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.view(head_mask.size(0), -1, 1, 1)
            else:
                hm = head_mask
            y = y * hm

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

        # KDA-FIX-02: output_scale via lora_output_scale (trainable)
        out = out * self.lora_output_scale
        if self.is_nested_encoder:
            final_out = out
        else:
            final_out = self.out_proj(out)

        return final_out, None

Every 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.

Methods

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
        if hasattr(mod, "freeze_base_weights"):
            mod.freeze_base_weights()

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class KernelFeatureConfig (head_dim,
num_features,
method='orthogonal',
clamp_min=-20.0,
clamp_max=8.0,
use_float64_exp=False,
eps=1e-05,
scaling=None,
seed=None)

Expand source code

class KernelFeatureConfig:
    def __init__(
        self,
        head_dim,
        num_features,
        method="orthogonal",
        clamp_min=-20.0,
        clamp_max=8.0,
        use_float64_exp=False,
        eps=1e-5,
        scaling=None,
        seed=None,
    ):
        self.head_dim = head_dim
        self.num_features = num_features
        self.method = method
        self.clamp_min = clamp_min
        self.clamp_max = clamp_max
        self.use_float64_exp = use_float64_exp
        self.eps = eps
        self.scaling = scaling
        self.seed = seed

class KernelFeatureMap (head_dim: int, num_features: int = 128, seed: int = 1234)

Expand source code

class KernelFeatureMap(nn.Module):
    """FAVOR+ style random Fourier features with deterministic seed."""

    is_custom_attention: bool = False

    def __init__(self, head_dim: int, num_features: int = 128, seed: int = 1234):
        super().__init__()
        self.head_dim = int(head_dim)
        self.num_features = int(num_features)
        gen = torch.Generator(device="cpu").manual_seed(int(seed))
        proj = torch.randn(self.head_dim, self.num_features, generator=gen)
        # Orthogonal init if m <= d
        if self.num_features <= self.head_dim:
            q, _ = torch.linalg.qr(torch.randn(self.head_dim, self.num_features, generator=gen))
            proj = q
        self.register_buffer("proj", proj * math.sqrt(self.head_dim))
        self.scale = 1.0 / math.sqrt(self.num_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.autocast(device_type=x.device.type, enabled=False):
            x_f = x.float()
            proj = self.proj.float()
            phi = torch.matmul(x_f, proj)
            norm_sq = (x_f * x_f).sum(dim=-1, keepdim=True) / 2.0
            raw = phi - norm_sq
            max_val = raw.detach().max()
            exp_in = torch.clamp(raw - max_val, min=-20.0, max=8.0)
            return (self.scale * torch.exp(exp_in)).to(x.dtype)

FAVOR+ style random Fourier features with deterministic seed.

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

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

Methods

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    with torch.autocast(device_type=x.device.type, enabled=False):
        x_f = x.float()
        proj = self.proj.float()
        phi = torch.matmul(x_f, proj)
        norm_sq = (x_f * x_f).sum(dim=-1, keepdim=True) / 2.0
        raw = phi - norm_sq
        max_val = raw.detach().max()
        exp_in = torch.clamp(raw - max_val, min=-20.0, max=8.0)
        return (self.scale * torch.exp(exp_in)).to(x.dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

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.0

Every 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-6

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).

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class LEKAV3Attention (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_v3", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("leka_v3_attention", supports_causal=True, supports_gqa=True)
class LEKAV3Attention(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

        # Robust super().__init__
        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_attention_heads = int(num_heads)
        self.num_heads = self.num_attention_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)

        # Hyper-parameters
        self.landmark_ratio = float(
            getattr(comp_config, "landmark_ratio", 0.25) if comp_config else 0.25
        )
        self.landmark_chunk_size = int(
            getattr(comp_config, "landmark_chunk_size", 64) if comp_config else 64
        )
        self.num_landmarks = int(
            kwargs.get("num_landmarks")
            or (getattr(comp_config, "num_landmarks", None) if comp_config else None)
            or (getattr(config, "num_landmarks", None) if config else None)
            or 64
        )
        self.feature_dim = int(getattr(comp_config, "feature_dim", 128) if comp_config else 128)
        self.local_window_size = int(
            kwargs.get("local_window_size")
            or (getattr(comp_config, "local_window_size", None) if comp_config else None)
            or (getattr(config, "local_window_size", None) if config else None)
            or 128
        )
        dropout_p = float(
            getattr(comp_config, "dropout", getattr(config, "dropout", 0.1))
            if config is not None
            else 0.1
        )

        # Centralized decoder detection
        self.is_decoder = _detect_decoder(config, is_decoder, source_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # LoRA config
        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)
        )

        # FactorizedLoRALinear projections
        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.dropout = nn.Dropout(dropout_p)

        # Single projection buffer (proj2 removed)
        self.register_buffer(
            "proj1",
            torch.randn(self.head_dim, self.feature_dim) / math.sqrt(self.head_dim / 2.0),
            persistent=True,
        )

        # Landmark scorer operates on K_features (2*feature_dim)
        self.landmark_scorer = nn.Sequential(
            nn.Linear(2 * self.feature_dim, 64),
            nn.GELU(),
            nn.Linear(64, 1),
        )

        # Per-head correlation matrix (identity init)
        self.W_corr = nn.Parameter(
            torch.eye(self.head_dim).unsqueeze(0).repeat(self.num_heads, 1, 1)
        )

        # Gate init = +2.0 → sigmoid(2) ≈ 0.88 favors LANDMARK branch
        self.gate = nn.Parameter(torch.full((self.num_heads,), 2.0))

        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", "query", "q_lin"):
                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))

    # ------------------------------------------------------------------
    # Kernel feature map (single projection)
    # ------------------------------------------------------------------
    def kernel_feature_map(self, x: torch.Tensor) -> torch.Tensor:
        proj1 = self.proj1.to(device=x.device, dtype=x.dtype)
        f1 = torch.matmul(x, proj1)
        f1 = torch.clamp(f1, min=-10.0, max=10.0)
        return torch.cat([torch.cos(f1), torch.sin(f1)], dim=-1) / math.sqrt(self.feature_dim)

    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
        new_x_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    # ------------------------------------------------------------------
    # Chunk-wise Landmark Selection (Solves Causal Starvation)
    # ------------------------------------------------------------------
    def select_landmarks_chunked(
        self,
        K_features: torch.Tensor,
        V: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
    ):
        B, H, N, feat_dim = K_features.shape
        chunk_sz = max(1, int(self.landmark_chunk_size))
        k_per_chunk = max(1, int(chunk_sz * self.landmark_ratio))

        with torch.no_grad():
            scores = self.landmark_scorer(K_features)[..., 0]  # (B, H, N)

            # Anchor boost for position 0 (always select first token)
            anchor_boost = torch.zeros_like(scores)
            anchor_boost[:, :, 0] = 1e4
            scores = scores + anchor_boost

            if valid_mask is not None:
                scores = scores.masked_fill(~valid_mask.unsqueeze(1).bool(), -1e4)

            # Pad to multiple of chunk_size
            pad_len = (chunk_sz - N % chunk_sz) % chunk_sz
            if pad_len > 0:
                scores_pad = F.pad(scores, (0, pad_len), value=-1e4)
            else:
                scores_pad = scores

            num_chunks = scores_pad.shape[-1] // chunk_sz
            scores_c = scores_pad.view(B, H, num_chunks, chunk_sz)

            # Extract top-K strictly within each temporal chunk
            k_act = min(k_per_chunk, chunk_sz)
            _, local_idx = scores_c.topk(k_act, dim=-1)

            # Convert local chunk indices back to global N-level indices
            offsets = torch.arange(0, num_chunks * chunk_sz, chunk_sz, device=scores.device).view(
                1, 1, -1, 1
            )
            global_idx = (local_idx + offsets).view(B, H, -1)
            global_idx = global_idx.clamp(max=N - 1)

        K_land = torch.gather(
            K_features,
            2,
            global_idx.unsqueeze(-1).expand(-1, -1, -1, feat_dim),
        )
        V_land = torch.gather(
            V,
            2,
            global_idx.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]),
        )
        return K_land, V_land, global_idx

    # ------------------------------------------------------------------
    # Block-wise Sliding Window Attention (True O(N) memory)
    # ------------------------------------------------------------------
    def block_local_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, N, d = q.shape
        W = max(1, int(self.local_window_size))

        pad_len = (W - N % W) % W
        if pad_len > 0:
            q_p = F.pad(q, (0, 0, 0, pad_len))
            k_p = F.pad(k, (0, 0, 0, pad_len))
            v_p = F.pad(v, (0, 0, 0, pad_len))
        else:
            q_p, k_p, v_p = q, k, v

        N_pad = q_p.shape[2]
        num_blocks = N_pad // W
        scale = 1.0 / math.sqrt(d)

        q_b = q_p.view(B, H, num_blocks, W, d)
        k_b = k_p.view(B, H, num_blocks, W, d)
        v_b = v_p.view(B, H, num_blocks, W, d)

        # 1. Intra-block (self) attention
        attn_intra = torch.matmul(q_b, k_b.transpose(-1, -2)) * scale
        if self.is_decoder:
            causal_mask = torch.tril(torch.ones(W, W, device=q.device, dtype=torch.bool))
            attn_intra = attn_intra.masked_fill(~causal_mask.view(1, 1, 1, W, W), -1e4)

        # 2. Inter-block (previous block) attention
        k_b_shifted = F.pad(k_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
        v_b_shifted = F.pad(v_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
        attn_inter = torch.matmul(q_b, k_b_shifted.transpose(-1, -2)) * scale
        attn_inter[:, :, 0, :, :] = -1e4  # First block has no previous block

        attn_combined = torch.cat([attn_inter, attn_intra], dim=-1)  # (B,H,nb,W,2W)

        if valid_mask is not None:
            v_mask_p = F.pad(valid_mask.float(), (0, pad_len), value=0.0).bool()
            # valid_mask shape: (B, N_pad) → broadcast against 2W dim
            v_mask_2w = torch.cat([v_mask_p, v_mask_p], dim=-1).view(B, num_blocks, 1, 2 * W)
            attn_combined = attn_combined.masked_fill(~v_mask_2w.unsqueeze(1), -1e4)

        attn_combined = torch.nan_to_num(attn_combined, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(attn_combined, dim=-1).to(v.dtype)
        attn_weights = self.dropout(attn_weights)

        v_combined = torch.cat([v_b_shifted, v_b], dim=-2)  # (B,H,nb,2W,d)
        out_local = torch.matmul(attn_weights, v_combined)  # (B,H,nb,W,d)

        return out_local.reshape(B, H, N_pad, d)[:, :, :N, :]

    # ------------------------------------------------------------------
    # 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("LEKAv3._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_()

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    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, LLaMA-style)
        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"LEKAV3Attention.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(
                            "LEKAV3Attention: 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()

        # Forward sanity check
        try:
            with torch.no_grad():
                d = next(self.parameters()).device
                dt = self.query.down.weight.dtype
                dummy = torch.randn(1, 8, self.dim, device=d, dtype=dt)
                _ = self.forward(dummy)
        except Exception as e:
            logger.debug("LEKAv3 forward sanity check warning: %s", e)

    # ------------------------------------------------------------------
    # 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 = self.num_heads
        d = self.head_dim
        device = hidden_states.device

        valid_mask = get_token_valid_mask(attention_mask, B, N, device)

        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()

        # RoPE with past_len offset
        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)

        Q_features = self.kernel_feature_map(q)
        K_features = self.kernel_feature_map(k)

        # 1. Landmark branch (chunk-wise selection, solves causal starvation)
        K_land, V_land, land_idx = self.select_landmarks_chunked(
            K_features,
            v,
            valid_mask=valid_mask,
        )

        # Per-head correlation via einsum
        # V_land: (B,H,L,d); W_corr: (H,d,d)
        # V_land_adj[b,h,l,e] = sum_d V_land[b,h,l,d] * W_corr[h,d,e]
        V_land_adj = torch.einsum(
            "bhld,hde->bhle",
            V_land,
            self.W_corr.to(V_land.dtype),
        )

        attn_land = torch.matmul(Q_features.float(), K_land.transpose(-2, -1).contiguous().float())

        if valid_mask is not None:
            pad_q = (~valid_mask).unsqueeze(1).unsqueeze(-1)
            attn_land = attn_land.masked_fill(pad_q.bool(), -1e4)

        if self.is_decoder:
            q_pos = torch.arange(N, device=device).view(1, 1, N, 1)
            causal_land_mask = land_idx.unsqueeze(2) <= q_pos
            attn_land = attn_land.masked_fill(~causal_land_mask, -1e4)

        attn_land = softpick(attn_land, dim=-1)
        attn_land = self.dropout(attn_land)
        context_land = torch.matmul(attn_land.to(v.dtype), V_land_adj)

        # 2. Local branch (block-wise sliding window, true O(N) memory)
        context_local = self.block_local_attention(q, k, v, valid_mask)

        # 3. Gate fusion: init sigmoid(+2.0)≈0.88 → favors LANDMARK branch
        gate_val = torch.sigmoid(self.gate).view(1, H, 1, 1).to(dtype=context_land.dtype)
        context = gate_val * context_land + (1.0 - gate_val) * context_local

        # 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)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return out, None

    def get_ortho_loss(self) -> float:
        return 0.0

Every 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 block_local_attention(self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
valid_mask: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def block_local_attention(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    B, H, N, d = q.shape
    W = max(1, int(self.local_window_size))

    pad_len = (W - N % W) % W
    if pad_len > 0:
        q_p = F.pad(q, (0, 0, 0, pad_len))
        k_p = F.pad(k, (0, 0, 0, pad_len))
        v_p = F.pad(v, (0, 0, 0, pad_len))
    else:
        q_p, k_p, v_p = q, k, v

    N_pad = q_p.shape[2]
    num_blocks = N_pad // W
    scale = 1.0 / math.sqrt(d)

    q_b = q_p.view(B, H, num_blocks, W, d)
    k_b = k_p.view(B, H, num_blocks, W, d)
    v_b = v_p.view(B, H, num_blocks, W, d)

    # 1. Intra-block (self) attention
    attn_intra = torch.matmul(q_b, k_b.transpose(-1, -2)) * scale
    if self.is_decoder:
        causal_mask = torch.tril(torch.ones(W, W, device=q.device, dtype=torch.bool))
        attn_intra = attn_intra.masked_fill(~causal_mask.view(1, 1, 1, W, W), -1e4)

    # 2. Inter-block (previous block) attention
    k_b_shifted = F.pad(k_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
    v_b_shifted = F.pad(v_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
    attn_inter = torch.matmul(q_b, k_b_shifted.transpose(-1, -2)) * scale
    attn_inter[:, :, 0, :, :] = -1e4  # First block has no previous block

    attn_combined = torch.cat([attn_inter, attn_intra], dim=-1)  # (B,H,nb,W,2W)

    if valid_mask is not None:
        v_mask_p = F.pad(valid_mask.float(), (0, pad_len), value=0.0).bool()
        # valid_mask shape: (B, N_pad) → broadcast against 2W dim
        v_mask_2w = torch.cat([v_mask_p, v_mask_p], dim=-1).view(B, num_blocks, 1, 2 * W)
        attn_combined = attn_combined.masked_fill(~v_mask_2w.unsqueeze(1), -1e4)

    attn_combined = torch.nan_to_num(attn_combined, nan=-1e4, posinf=1e4, neginf=-1e4)
    attn_weights = F.softmax(attn_combined, dim=-1).to(v.dtype)
    attn_weights = self.dropout(attn_weights)

    v_combined = torch.cat([v_b_shifted, v_b], dim=-2)  # (B,H,nb,2W,d)
    out_local = torch.matmul(attn_weights, v_combined)  # (B,H,nb,W,d)

    return out_local.reshape(B, H, N_pad, d)[:, :, :N, :]

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

Expand source code

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, LLaMA-style)
    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"LEKAV3Attention.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(
                        "LEKAV3Attention: 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()

    # Forward sanity check
    try:
        with torch.no_grad():
            d = next(self.parameters()).device
            dt = self.query.down.weight.dtype
            dummy = torch.randn(1, 8, self.dim, device=d, dtype=dt)
            _ = self.forward(dummy)
    except Exception as e:
        logger.debug("LEKAv3 forward sanity check warning: %s", e)

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,
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]

Expand source code

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 = self.num_heads
    d = self.head_dim
    device = hidden_states.device

    valid_mask = get_token_valid_mask(attention_mask, B, N, device)

    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()

    # RoPE with past_len offset
    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)

    Q_features = self.kernel_feature_map(q)
    K_features = self.kernel_feature_map(k)

    # 1. Landmark branch (chunk-wise selection, solves causal starvation)
    K_land, V_land, land_idx = self.select_landmarks_chunked(
        K_features,
        v,
        valid_mask=valid_mask,
    )

    # Per-head correlation via einsum
    # V_land: (B,H,L,d); W_corr: (H,d,d)
    # V_land_adj[b,h,l,e] = sum_d V_land[b,h,l,d] * W_corr[h,d,e]
    V_land_adj = torch.einsum(
        "bhld,hde->bhle",
        V_land,
        self.W_corr.to(V_land.dtype),
    )

    attn_land = torch.matmul(Q_features.float(), K_land.transpose(-2, -1).contiguous().float())

    if valid_mask is not None:
        pad_q = (~valid_mask).unsqueeze(1).unsqueeze(-1)
        attn_land = attn_land.masked_fill(pad_q.bool(), -1e4)

    if self.is_decoder:
        q_pos = torch.arange(N, device=device).view(1, 1, N, 1)
        causal_land_mask = land_idx.unsqueeze(2) <= q_pos
        attn_land = attn_land.masked_fill(~causal_land_mask, -1e4)

    attn_land = softpick(attn_land, dim=-1)
    attn_land = self.dropout(attn_land)
    context_land = torch.matmul(attn_land.to(v.dtype), V_land_adj)

    # 2. Local branch (block-wise sliding window, true O(N) memory)
    context_local = self.block_local_attention(q, k, v, valid_mask)

    # 3. Gate fusion: init sigmoid(+2.0)≈0.88 → favors LANDMARK branch
    gate_val = torch.sigmoid(self.gate).view(1, H, 1, 1).to(dtype=context_land.dtype)
    context = gate_val * context_land + (1.0 - gate_val) * context_local

    # 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)
    context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

    return out, None

Must return (output, attention_weights_or_None).

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:
    proj1 = self.proj1.to(device=x.device, dtype=x.dtype)
    f1 = torch.matmul(x, proj1)
    f1 = torch.clamp(f1, min=-10.0, max=10.0)
    return torch.cat([torch.cos(f1), torch.sin(f1)], dim=-1) / math.sqrt(self.feature_dim)

def select_landmarks_chunked(self,
K_features: torch.Tensor,
V: torch.Tensor,
valid_mask: torch.Tensor | None = None)

Expand source code

def select_landmarks_chunked(
    self,
    K_features: torch.Tensor,
    V: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
):
    B, H, N, feat_dim = K_features.shape
    chunk_sz = max(1, int(self.landmark_chunk_size))
    k_per_chunk = max(1, int(chunk_sz * self.landmark_ratio))

    with torch.no_grad():
        scores = self.landmark_scorer(K_features)[..., 0]  # (B, H, N)

        # Anchor boost for position 0 (always select first token)
        anchor_boost = torch.zeros_like(scores)
        anchor_boost[:, :, 0] = 1e4
        scores = scores + anchor_boost

        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask.unsqueeze(1).bool(), -1e4)

        # Pad to multiple of chunk_size
        pad_len = (chunk_sz - N % chunk_sz) % chunk_sz
        if pad_len > 0:
            scores_pad = F.pad(scores, (0, pad_len), value=-1e4)
        else:
            scores_pad = scores

        num_chunks = scores_pad.shape[-1] // chunk_sz
        scores_c = scores_pad.view(B, H, num_chunks, chunk_sz)

        # Extract top-K strictly within each temporal chunk
        k_act = min(k_per_chunk, chunk_sz)
        _, local_idx = scores_c.topk(k_act, dim=-1)

        # Convert local chunk indices back to global N-level indices
        offsets = torch.arange(0, num_chunks * chunk_sz, chunk_sz, device=scores.device).view(
            1, 1, -1, 1
        )
        global_idx = (local_idx + offsets).view(B, H, -1)
        global_idx = global_idx.clamp(max=N - 1)

    K_land = torch.gather(
        K_features,
        2,
        global_idx.unsqueeze(-1).expand(-1, -1, -1, feat_dim),
    )
    V_land = torch.gather(
        V,
        2,
        global_idx.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]),
    )
    return K_land, V_land, global_idx

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

Expand source code

def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
    new_x_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
    x = x.view(new_x_shape)
    return x.permute(0, 2, 1, 3)

class LEKAv3Attention (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_v3", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("leka_v3_attention", supports_causal=True, supports_gqa=True)
class LEKAV3Attention(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

        # Robust super().__init__
        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_attention_heads = int(num_heads)
        self.num_heads = self.num_attention_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)

        # Hyper-parameters
        self.landmark_ratio = float(
            getattr(comp_config, "landmark_ratio", 0.25) if comp_config else 0.25
        )
        self.landmark_chunk_size = int(
            getattr(comp_config, "landmark_chunk_size", 64) if comp_config else 64
        )
        self.num_landmarks = int(
            kwargs.get("num_landmarks")
            or (getattr(comp_config, "num_landmarks", None) if comp_config else None)
            or (getattr(config, "num_landmarks", None) if config else None)
            or 64
        )
        self.feature_dim = int(getattr(comp_config, "feature_dim", 128) if comp_config else 128)
        self.local_window_size = int(
            kwargs.get("local_window_size")
            or (getattr(comp_config, "local_window_size", None) if comp_config else None)
            or (getattr(config, "local_window_size", None) if config else None)
            or 128
        )
        dropout_p = float(
            getattr(comp_config, "dropout", getattr(config, "dropout", 0.1))
            if config is not None
            else 0.1
        )

        # Centralized decoder detection
        self.is_decoder = _detect_decoder(config, is_decoder, source_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # LoRA config
        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)
        )

        # FactorizedLoRALinear projections
        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.dropout = nn.Dropout(dropout_p)

        # Single projection buffer (proj2 removed)
        self.register_buffer(
            "proj1",
            torch.randn(self.head_dim, self.feature_dim) / math.sqrt(self.head_dim / 2.0),
            persistent=True,
        )

        # Landmark scorer operates on K_features (2*feature_dim)
        self.landmark_scorer = nn.Sequential(
            nn.Linear(2 * self.feature_dim, 64),
            nn.GELU(),
            nn.Linear(64, 1),
        )

        # Per-head correlation matrix (identity init)
        self.W_corr = nn.Parameter(
            torch.eye(self.head_dim).unsqueeze(0).repeat(self.num_heads, 1, 1)
        )

        # Gate init = +2.0 → sigmoid(2) ≈ 0.88 favors LANDMARK branch
        self.gate = nn.Parameter(torch.full((self.num_heads,), 2.0))

        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", "query", "q_lin"):
                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))

    # ------------------------------------------------------------------
    # Kernel feature map (single projection)
    # ------------------------------------------------------------------
    def kernel_feature_map(self, x: torch.Tensor) -> torch.Tensor:
        proj1 = self.proj1.to(device=x.device, dtype=x.dtype)
        f1 = torch.matmul(x, proj1)
        f1 = torch.clamp(f1, min=-10.0, max=10.0)
        return torch.cat([torch.cos(f1), torch.sin(f1)], dim=-1) / math.sqrt(self.feature_dim)

    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
        new_x_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    # ------------------------------------------------------------------
    # Chunk-wise Landmark Selection (Solves Causal Starvation)
    # ------------------------------------------------------------------
    def select_landmarks_chunked(
        self,
        K_features: torch.Tensor,
        V: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
    ):
        B, H, N, feat_dim = K_features.shape
        chunk_sz = max(1, int(self.landmark_chunk_size))
        k_per_chunk = max(1, int(chunk_sz * self.landmark_ratio))

        with torch.no_grad():
            scores = self.landmark_scorer(K_features)[..., 0]  # (B, H, N)

            # Anchor boost for position 0 (always select first token)
            anchor_boost = torch.zeros_like(scores)
            anchor_boost[:, :, 0] = 1e4
            scores = scores + anchor_boost

            if valid_mask is not None:
                scores = scores.masked_fill(~valid_mask.unsqueeze(1).bool(), -1e4)

            # Pad to multiple of chunk_size
            pad_len = (chunk_sz - N % chunk_sz) % chunk_sz
            if pad_len > 0:
                scores_pad = F.pad(scores, (0, pad_len), value=-1e4)
            else:
                scores_pad = scores

            num_chunks = scores_pad.shape[-1] // chunk_sz
            scores_c = scores_pad.view(B, H, num_chunks, chunk_sz)

            # Extract top-K strictly within each temporal chunk
            k_act = min(k_per_chunk, chunk_sz)
            _, local_idx = scores_c.topk(k_act, dim=-1)

            # Convert local chunk indices back to global N-level indices
            offsets = torch.arange(0, num_chunks * chunk_sz, chunk_sz, device=scores.device).view(
                1, 1, -1, 1
            )
            global_idx = (local_idx + offsets).view(B, H, -1)
            global_idx = global_idx.clamp(max=N - 1)

        K_land = torch.gather(
            K_features,
            2,
            global_idx.unsqueeze(-1).expand(-1, -1, -1, feat_dim),
        )
        V_land = torch.gather(
            V,
            2,
            global_idx.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]),
        )
        return K_land, V_land, global_idx

    # ------------------------------------------------------------------
    # Block-wise Sliding Window Attention (True O(N) memory)
    # ------------------------------------------------------------------
    def block_local_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, N, d = q.shape
        W = max(1, int(self.local_window_size))

        pad_len = (W - N % W) % W
        if pad_len > 0:
            q_p = F.pad(q, (0, 0, 0, pad_len))
            k_p = F.pad(k, (0, 0, 0, pad_len))
            v_p = F.pad(v, (0, 0, 0, pad_len))
        else:
            q_p, k_p, v_p = q, k, v

        N_pad = q_p.shape[2]
        num_blocks = N_pad // W
        scale = 1.0 / math.sqrt(d)

        q_b = q_p.view(B, H, num_blocks, W, d)
        k_b = k_p.view(B, H, num_blocks, W, d)
        v_b = v_p.view(B, H, num_blocks, W, d)

        # 1. Intra-block (self) attention
        attn_intra = torch.matmul(q_b, k_b.transpose(-1, -2)) * scale
        if self.is_decoder:
            causal_mask = torch.tril(torch.ones(W, W, device=q.device, dtype=torch.bool))
            attn_intra = attn_intra.masked_fill(~causal_mask.view(1, 1, 1, W, W), -1e4)

        # 2. Inter-block (previous block) attention
        k_b_shifted = F.pad(k_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
        v_b_shifted = F.pad(v_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
        attn_inter = torch.matmul(q_b, k_b_shifted.transpose(-1, -2)) * scale
        attn_inter[:, :, 0, :, :] = -1e4  # First block has no previous block

        attn_combined = torch.cat([attn_inter, attn_intra], dim=-1)  # (B,H,nb,W,2W)

        if valid_mask is not None:
            v_mask_p = F.pad(valid_mask.float(), (0, pad_len), value=0.0).bool()
            # valid_mask shape: (B, N_pad) → broadcast against 2W dim
            v_mask_2w = torch.cat([v_mask_p, v_mask_p], dim=-1).view(B, num_blocks, 1, 2 * W)
            attn_combined = attn_combined.masked_fill(~v_mask_2w.unsqueeze(1), -1e4)

        attn_combined = torch.nan_to_num(attn_combined, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(attn_combined, dim=-1).to(v.dtype)
        attn_weights = self.dropout(attn_weights)

        v_combined = torch.cat([v_b_shifted, v_b], dim=-2)  # (B,H,nb,2W,d)
        out_local = torch.matmul(attn_weights, v_combined)  # (B,H,nb,W,d)

        return out_local.reshape(B, H, N_pad, d)[:, :, :N, :]

    # ------------------------------------------------------------------
    # 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("LEKAv3._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_()

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    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, LLaMA-style)
        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"LEKAV3Attention.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(
                            "LEKAV3Attention: 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()

        # Forward sanity check
        try:
            with torch.no_grad():
                d = next(self.parameters()).device
                dt = self.query.down.weight.dtype
                dummy = torch.randn(1, 8, self.dim, device=d, dtype=dt)
                _ = self.forward(dummy)
        except Exception as e:
            logger.debug("LEKAv3 forward sanity check warning: %s", e)

    # ------------------------------------------------------------------
    # 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 = self.num_heads
        d = self.head_dim
        device = hidden_states.device

        valid_mask = get_token_valid_mask(attention_mask, B, N, device)

        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()

        # RoPE with past_len offset
        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)

        Q_features = self.kernel_feature_map(q)
        K_features = self.kernel_feature_map(k)

        # 1. Landmark branch (chunk-wise selection, solves causal starvation)
        K_land, V_land, land_idx = self.select_landmarks_chunked(
            K_features,
            v,
            valid_mask=valid_mask,
        )

        # Per-head correlation via einsum
        # V_land: (B,H,L,d); W_corr: (H,d,d)
        # V_land_adj[b,h,l,e] = sum_d V_land[b,h,l,d] * W_corr[h,d,e]
        V_land_adj = torch.einsum(
            "bhld,hde->bhle",
            V_land,
            self.W_corr.to(V_land.dtype),
        )

        attn_land = torch.matmul(Q_features.float(), K_land.transpose(-2, -1).contiguous().float())

        if valid_mask is not None:
            pad_q = (~valid_mask).unsqueeze(1).unsqueeze(-1)
            attn_land = attn_land.masked_fill(pad_q.bool(), -1e4)

        if self.is_decoder:
            q_pos = torch.arange(N, device=device).view(1, 1, N, 1)
            causal_land_mask = land_idx.unsqueeze(2) <= q_pos
            attn_land = attn_land.masked_fill(~causal_land_mask, -1e4)

        attn_land = softpick(attn_land, dim=-1)
        attn_land = self.dropout(attn_land)
        context_land = torch.matmul(attn_land.to(v.dtype), V_land_adj)

        # 2. Local branch (block-wise sliding window, true O(N) memory)
        context_local = self.block_local_attention(q, k, v, valid_mask)

        # 3. Gate fusion: init sigmoid(+2.0)≈0.88 → favors LANDMARK branch
        gate_val = torch.sigmoid(self.gate).view(1, H, 1, 1).to(dtype=context_land.dtype)
        context = gate_val * context_land + (1.0 - gate_val) * context_local

        # 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)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return out, None

    def get_ortho_loss(self) -> float:
        return 0.0

Every 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 block_local_attention(self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
valid_mask: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def block_local_attention(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    B, H, N, d = q.shape
    W = max(1, int(self.local_window_size))

    pad_len = (W - N % W) % W
    if pad_len > 0:
        q_p = F.pad(q, (0, 0, 0, pad_len))
        k_p = F.pad(k, (0, 0, 0, pad_len))
        v_p = F.pad(v, (0, 0, 0, pad_len))
    else:
        q_p, k_p, v_p = q, k, v

    N_pad = q_p.shape[2]
    num_blocks = N_pad // W
    scale = 1.0 / math.sqrt(d)

    q_b = q_p.view(B, H, num_blocks, W, d)
    k_b = k_p.view(B, H, num_blocks, W, d)
    v_b = v_p.view(B, H, num_blocks, W, d)

    # 1. Intra-block (self) attention
    attn_intra = torch.matmul(q_b, k_b.transpose(-1, -2)) * scale
    if self.is_decoder:
        causal_mask = torch.tril(torch.ones(W, W, device=q.device, dtype=torch.bool))
        attn_intra = attn_intra.masked_fill(~causal_mask.view(1, 1, 1, W, W), -1e4)

    # 2. Inter-block (previous block) attention
    k_b_shifted = F.pad(k_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
    v_b_shifted = F.pad(v_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
    attn_inter = torch.matmul(q_b, k_b_shifted.transpose(-1, -2)) * scale
    attn_inter[:, :, 0, :, :] = -1e4  # First block has no previous block

    attn_combined = torch.cat([attn_inter, attn_intra], dim=-1)  # (B,H,nb,W,2W)

    if valid_mask is not None:
        v_mask_p = F.pad(valid_mask.float(), (0, pad_len), value=0.0).bool()
        # valid_mask shape: (B, N_pad) → broadcast against 2W dim
        v_mask_2w = torch.cat([v_mask_p, v_mask_p], dim=-1).view(B, num_blocks, 1, 2 * W)
        attn_combined = attn_combined.masked_fill(~v_mask_2w.unsqueeze(1), -1e4)

    attn_combined = torch.nan_to_num(attn_combined, nan=-1e4, posinf=1e4, neginf=-1e4)
    attn_weights = F.softmax(attn_combined, dim=-1).to(v.dtype)
    attn_weights = self.dropout(attn_weights)

    v_combined = torch.cat([v_b_shifted, v_b], dim=-2)  # (B,H,nb,2W,d)
    out_local = torch.matmul(attn_weights, v_combined)  # (B,H,nb,W,d)

    return out_local.reshape(B, H, N_pad, d)[:, :, :N, :]

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:
    proj1 = self.proj1.to(device=x.device, dtype=x.dtype)
    f1 = torch.matmul(x, proj1)
    f1 = torch.clamp(f1, min=-10.0, max=10.0)
    return torch.cat([torch.cos(f1), torch.sin(f1)], dim=-1) / math.sqrt(self.feature_dim)

def select_landmarks_chunked(self,
K_features: torch.Tensor,
V: torch.Tensor,
valid_mask: torch.Tensor | None = None)

Expand source code

def select_landmarks_chunked(
    self,
    K_features: torch.Tensor,
    V: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
):
    B, H, N, feat_dim = K_features.shape
    chunk_sz = max(1, int(self.landmark_chunk_size))
    k_per_chunk = max(1, int(chunk_sz * self.landmark_ratio))

    with torch.no_grad():
        scores = self.landmark_scorer(K_features)[..., 0]  # (B, H, N)

        # Anchor boost for position 0 (always select first token)
        anchor_boost = torch.zeros_like(scores)
        anchor_boost[:, :, 0] = 1e4
        scores = scores + anchor_boost

        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask.unsqueeze(1).bool(), -1e4)

        # Pad to multiple of chunk_size
        pad_len = (chunk_sz - N % chunk_sz) % chunk_sz
        if pad_len > 0:
            scores_pad = F.pad(scores, (0, pad_len), value=-1e4)
        else:
            scores_pad = scores

        num_chunks = scores_pad.shape[-1] // chunk_sz
        scores_c = scores_pad.view(B, H, num_chunks, chunk_sz)

        # Extract top-K strictly within each temporal chunk
        k_act = min(k_per_chunk, chunk_sz)
        _, local_idx = scores_c.topk(k_act, dim=-1)

        # Convert local chunk indices back to global N-level indices
        offsets = torch.arange(0, num_chunks * chunk_sz, chunk_sz, device=scores.device).view(
            1, 1, -1, 1
        )
        global_idx = (local_idx + offsets).view(B, H, -1)
        global_idx = global_idx.clamp(max=N - 1)

    K_land = torch.gather(
        K_features,
        2,
        global_idx.unsqueeze(-1).expand(-1, -1, -1, feat_dim),
    )
    V_land = torch.gather(
        V,
        2,
        global_idx.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]),
    )
    return K_land, V_land, global_idx

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

Expand source code

def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
    new_x_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
    x = x.view(new_x_shape)
    return x.permute(0, 2, 1, 3)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class LEKAV2Attention (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_v2", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("leka_v2_attention", supports_causal=True, supports_gqa=True)
class LEKAv2Attention(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_attention_heads = int(num_heads)
        self.num_heads = self.num_attention_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)

        # Hyper-parameters
        self.landmark_ratio = float(
            getattr(comp_config, "landmark_ratio", 0.5) if comp_config else 0.5
        )
        self.landmark_chunk_size = int(
            getattr(comp_config, "landmark_chunk_size", 64) if comp_config else 64
        )
        self.min_landmarks = int(getattr(comp_config, "min_landmarks", 64) if comp_config else 64)
        self.num_landmarks = int(
            getattr(comp_config, "num_landmarks", self.min_landmarks)
            if comp_config
            else self.min_landmarks
        )
        self.feature_dim = int(getattr(comp_config, "feature_dim", 128) if comp_config else 128)
        self.local_window_size = int(
            getattr(comp_config, "local_window_size", 128) if comp_config else 128
        )
        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.k_proj, self.v_proj = self.query, self.key, self.value

        self.register_buffer(
            "rand_proj1",
            torch.randn(self.head_dim, self.feature_dim) / math.sqrt(self.head_dim),
            persistent=True,
        )
        self.register_buffer(
            "rand_proj2",
            torch.randn(self.head_dim, self.feature_dim // 2) / math.sqrt(self.head_dim),
            persistent=True,
        )

        tot_feat = self.feature_dim * 2 + (self.feature_dim // 2) * 2
        self.landmark_scorer = nn.Sequential(
            nn.LayerNorm(tot_feat),
            nn.Linear(tot_feat, 64),
            nn.GELU(),
            nn.Linear(64, 1),
        )

        self.W_corr = nn.Parameter(
            torch.eye(self.head_dim).unsqueeze(0).repeat(self.num_heads, 1, 1)
        )
        self.local_gate = nn.Parameter(torch.full((self.num_heads,), -2.0))
        self.dropout = nn.Dropout(dropout_p)

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

    # ------------------------------------------------------------------
    # Dimension detectors
    # ------------------------------------------------------------------
    @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", "query", "q_lin"):
                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))

    # ------------------------------------------------------------------
    # Kernel feature map (dual-scale RFF)
    # ------------------------------------------------------------------
    def kernel_feature_map(self, x: torch.Tensor) -> torch.Tensor:
        W1 = self.rand_proj1.to(device=x.device, dtype=x.dtype)
        W2 = self.rand_proj2.to(device=x.device, dtype=x.dtype)

        proj1 = torch.matmul(x, W1)
        f1 = torch.cat([torch.cos(proj1), torch.sin(proj1)], dim=-1) / math.sqrt(self.feature_dim)

        proj2 = torch.matmul(x, W2)
        f2 = torch.cat([torch.cos(proj2), torch.sin(proj2)], dim=-1) / math.sqrt(
            self.feature_dim // 2
        )

        return torch.cat([f1, f2], dim=-1)

    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
        new_x_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    # ------------------------------------------------------------------
    # Chunk-wise Landmark Selection (Solves Causal Starvation)
    # ------------------------------------------------------------------
    def select_landmarks_chunked(
        self,
        K_features: torch.Tensor,
        V: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
    ):
        B, H, N, feat_dim = K_features.shape
        chunk_sz = max(1, int(self.landmark_chunk_size))
        k_per_chunk = max(1, int(chunk_sz * self.landmark_ratio))

        with torch.no_grad():
            scores = self.landmark_scorer(K_features).squeeze(-1)  # (B, H, N)

            # Anchor boost for position 0 (always select the first token)
            anchor_boost = torch.zeros_like(scores)
            anchor_boost[:, :, 0] = 1e4
            scores = scores + anchor_boost

            if valid_mask is not None:
                scores = scores.masked_fill(~valid_mask.unsqueeze(1), -1e4)

            # Pad to multiple of chunk_size
            pad_len = (chunk_sz - N % chunk_sz) % chunk_sz
            if pad_len > 0:
                scores_pad = F.pad(scores, (0, pad_len), value=-1e4)
            else:
                scores_pad = scores

            num_chunks = scores_pad.shape[-1] // chunk_sz
            scores_c = scores_pad.view(B, H, num_chunks, chunk_sz)

            k_act = min(k_per_chunk, chunk_sz)
            _, local_idx = scores_c.topk(k_act, dim=-1)

            offsets = torch.arange(0, num_chunks * chunk_sz, chunk_sz, device=scores.device).view(
                1, 1, -1, 1
            )
            global_idx = (local_idx + offsets).view(B, H, -1)
            global_idx = global_idx.clamp(max=N - 1)

        K_land = torch.gather(
            K_features,
            2,
            global_idx.unsqueeze(-1).expand(-1, -1, -1, feat_dim),
        )
        V_land = torch.gather(
            V,
            2,
            global_idx.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]),
        )
        return K_land, V_land, global_idx

    # ------------------------------------------------------------------
    # Block-wise Sliding Window Attention (True O(N) memory)
    # ------------------------------------------------------------------
    def block_local_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, N, d = q.shape
        W = max(1, int(self.local_window_size))

        pad_len = (W - N % W) % W
        if pad_len > 0:
            q_p = F.pad(q, (0, 0, 0, pad_len))
            k_p = F.pad(k, (0, 0, 0, pad_len))
            v_p = F.pad(v, (0, 0, 0, pad_len))
        else:
            q_p, k_p, v_p = q, k, v

        N_pad = q_p.shape[2]
        num_blocks = N_pad // W
        scale = 1.0 / math.sqrt(d)

        q_b = q_p.view(B, H, num_blocks, W, d)
        k_b = k_p.view(B, H, num_blocks, W, d)
        v_b = v_p.view(B, H, num_blocks, W, d)

        # 1. Intra-block (self) attention
        attn_intra = torch.matmul(q_b, k_b.transpose(-1, -2)) * scale
        if self.is_decoder:
            causal_mask = torch.tril(torch.ones(W, W, device=q.device, dtype=torch.bool))
            attn_intra = attn_intra.masked_fill(~causal_mask.view(1, 1, 1, W, W), -1e4)

        # 2. Inter-block (previous block) attention
        k_b_shifted = F.pad(k_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
        v_b_shifted = F.pad(v_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
        attn_inter = torch.matmul(q_b, k_b_shifted.transpose(-1, -2)) * scale
        attn_inter[:, :, 0, :, :] = -1e4  # First block has no previous block

        attn_combined = torch.cat([attn_inter, attn_intra], dim=-1)  # (B, H, nb, W, 2W)

        if valid_mask is not None:
            v_mask_p = F.pad(valid_mask.float(), (0, pad_len), value=0.0).bool()
            # valid_mask shape: (B, N_pad) → broadcast against 2W dim
            v_mask_2w = torch.cat([v_mask_p, v_mask_p], dim=-1).view(B, num_blocks, 1, 2 * W)
            attn_combined = attn_combined.masked_fill(~v_mask_2w.unsqueeze(1), -1e4)

        attn_combined = torch.nan_to_num(attn_combined, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(attn_combined, dim=-1).to(v.dtype)
        attn_weights = self.dropout(attn_weights)

        v_combined = torch.cat([v_b_shifted, v_b], dim=-2)  # (B, H, nb, 2W, d)
        out_local = torch.matmul(attn_weights, v_combined)  # (B, H, nb, W, d)

        return out_local.reshape(B, H, N_pad, d)[:, :, :N, :]

    # ------------------------------------------------------------------
    # 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("LEKAv2._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_()

    # ------------------------------------------------------------------
    # copy_weights (GPT-2 / BERT-nested / DistilBERT / Qwen)
    # ------------------------------------------------------------------
    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, LLaMA-style)
        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"LEKAv2Attention.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(
                            "LEKAv2Attention: 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()

        # Forward sanity check
        try:
            with torch.no_grad():
                d = next(self.parameters()).device
                dt = self.query.down.weight.dtype
                dummy = torch.randn(1, 8, self.dim, device=d, dtype=dt)
                _ = self.forward(dummy)
        except Exception as e:
            logger.debug("LEKAv2 forward sanity check warning: %s", e)

    # ------------------------------------------------------------------
    # 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 = self.num_heads
        device = hidden_states.device

        q = self.query(hidden_states).view(B, N, H, self.head_dim).transpose(1, 2).contiguous()
        k = (
            self.key(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.value(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # RoPE with past_len offset
        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)

        Q_features = self.kernel_feature_map(q)
        K_features = self.kernel_feature_map(k)

        token_valid = None
        if attention_mask is not None:
            token_valid = get_token_valid_mask(attention_mask, B, N, device)

        # 1. Landmark branch (chunk-wise selection, solves causal starvation)
        K_land, V_land, land_idx = self.select_landmarks_chunked(
            K_features,
            v,
            valid_mask=token_valid,
        )
        V_land_adj = torch.einsum("bhld,hde->bhle", V_land, self.W_corr.to(V_land.dtype))

        attn_land = torch.matmul(Q_features.float(), K_land.transpose(-2, -1).contiguous().float())

        if token_valid is not None:
            attn_land = attn_land.masked_fill(~token_valid.unsqueeze(1).unsqueeze(-1), -1e4)

        if self.is_decoder:
            q_pos = torch.arange(N, device=device).view(1, 1, N, 1)
            causal_land_mask = land_idx.unsqueeze(2) <= q_pos
            attn_land = attn_land.masked_fill(~causal_land_mask, -1e4)

        attn_land = softpick(attn_land, dim=-1)
        attn_land = self.dropout(attn_land)
        context_leka = torch.matmul(attn_land.to(v.dtype), V_land_adj)

        # 2. Local branch (block-wise sliding window, true O(N) memory)
        context_local = self.block_local_attention(q, k, v, token_valid)

        # 3. Gate fusion (local dominates at init)
        gate_val = torch.sigmoid(self.local_gate).view(1, H, 1, 1)
        context = (1.0 - gate_val) * context_local + gate_val * context_leka

        # 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)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return out, None

    def get_ortho_loss(self) -> float:
        return 0.0

Every 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 block_local_attention(self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
valid_mask: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def block_local_attention(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    B, H, N, d = q.shape
    W = max(1, int(self.local_window_size))

    pad_len = (W - N % W) % W
    if pad_len > 0:
        q_p = F.pad(q, (0, 0, 0, pad_len))
        k_p = F.pad(k, (0, 0, 0, pad_len))
        v_p = F.pad(v, (0, 0, 0, pad_len))
    else:
        q_p, k_p, v_p = q, k, v

    N_pad = q_p.shape[2]
    num_blocks = N_pad // W
    scale = 1.0 / math.sqrt(d)

    q_b = q_p.view(B, H, num_blocks, W, d)
    k_b = k_p.view(B, H, num_blocks, W, d)
    v_b = v_p.view(B, H, num_blocks, W, d)

    # 1. Intra-block (self) attention
    attn_intra = torch.matmul(q_b, k_b.transpose(-1, -2)) * scale
    if self.is_decoder:
        causal_mask = torch.tril(torch.ones(W, W, device=q.device, dtype=torch.bool))
        attn_intra = attn_intra.masked_fill(~causal_mask.view(1, 1, 1, W, W), -1e4)

    # 2. Inter-block (previous block) attention
    k_b_shifted = F.pad(k_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
    v_b_shifted = F.pad(v_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
    attn_inter = torch.matmul(q_b, k_b_shifted.transpose(-1, -2)) * scale
    attn_inter[:, :, 0, :, :] = -1e4  # First block has no previous block

    attn_combined = torch.cat([attn_inter, attn_intra], dim=-1)  # (B, H, nb, W, 2W)

    if valid_mask is not None:
        v_mask_p = F.pad(valid_mask.float(), (0, pad_len), value=0.0).bool()
        # valid_mask shape: (B, N_pad) → broadcast against 2W dim
        v_mask_2w = torch.cat([v_mask_p, v_mask_p], dim=-1).view(B, num_blocks, 1, 2 * W)
        attn_combined = attn_combined.masked_fill(~v_mask_2w.unsqueeze(1), -1e4)

    attn_combined = torch.nan_to_num(attn_combined, nan=-1e4, posinf=1e4, neginf=-1e4)
    attn_weights = F.softmax(attn_combined, dim=-1).to(v.dtype)
    attn_weights = self.dropout(attn_weights)

    v_combined = torch.cat([v_b_shifted, v_b], dim=-2)  # (B, H, nb, 2W, d)
    out_local = torch.matmul(attn_weights, v_combined)  # (B, H, nb, W, d)

    return out_local.reshape(B, H, N_pad, d)[:, :, :N, :]

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

Expand source code

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, LLaMA-style)
    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"LEKAv2Attention.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(
                        "LEKAv2Attention: 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()

    # Forward sanity check
    try:
        with torch.no_grad():
            d = next(self.parameters()).device
            dt = self.query.down.weight.dtype
            dummy = torch.randn(1, 8, self.dim, device=d, dtype=dt)
            _ = self.forward(dummy)
    except Exception as e:
        logger.debug("LEKAv2 forward sanity check warning: %s", e)

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,
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]

Expand source code

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 = self.num_heads
    device = hidden_states.device

    q = self.query(hidden_states).view(B, N, H, self.head_dim).transpose(1, 2).contiguous()
    k = (
        self.key(hidden_states)
        .view(B, N, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v = (
        self.value(hidden_states)
        .view(B, N, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # RoPE with past_len offset
    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)

    Q_features = self.kernel_feature_map(q)
    K_features = self.kernel_feature_map(k)

    token_valid = None
    if attention_mask is not None:
        token_valid = get_token_valid_mask(attention_mask, B, N, device)

    # 1. Landmark branch (chunk-wise selection, solves causal starvation)
    K_land, V_land, land_idx = self.select_landmarks_chunked(
        K_features,
        v,
        valid_mask=token_valid,
    )
    V_land_adj = torch.einsum("bhld,hde->bhle", V_land, self.W_corr.to(V_land.dtype))

    attn_land = torch.matmul(Q_features.float(), K_land.transpose(-2, -1).contiguous().float())

    if token_valid is not None:
        attn_land = attn_land.masked_fill(~token_valid.unsqueeze(1).unsqueeze(-1), -1e4)

    if self.is_decoder:
        q_pos = torch.arange(N, device=device).view(1, 1, N, 1)
        causal_land_mask = land_idx.unsqueeze(2) <= q_pos
        attn_land = attn_land.masked_fill(~causal_land_mask, -1e4)

    attn_land = softpick(attn_land, dim=-1)
    attn_land = self.dropout(attn_land)
    context_leka = torch.matmul(attn_land.to(v.dtype), V_land_adj)

    # 2. Local branch (block-wise sliding window, true O(N) memory)
    context_local = self.block_local_attention(q, k, v, token_valid)

    # 3. Gate fusion (local dominates at init)
    gate_val = torch.sigmoid(self.local_gate).view(1, H, 1, 1)
    context = (1.0 - gate_val) * context_local + gate_val * context_leka

    # 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)
    context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

    return out, None

Must return (output, attention_weights_or_None).

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:
    W1 = self.rand_proj1.to(device=x.device, dtype=x.dtype)
    W2 = self.rand_proj2.to(device=x.device, dtype=x.dtype)

    proj1 = torch.matmul(x, W1)
    f1 = torch.cat([torch.cos(proj1), torch.sin(proj1)], dim=-1) / math.sqrt(self.feature_dim)

    proj2 = torch.matmul(x, W2)
    f2 = torch.cat([torch.cos(proj2), torch.sin(proj2)], dim=-1) / math.sqrt(
        self.feature_dim // 2
    )

    return torch.cat([f1, f2], dim=-1)

def select_landmarks_chunked(self,
K_features: torch.Tensor,
V: torch.Tensor,
valid_mask: torch.Tensor | None = None)

Expand source code

def select_landmarks_chunked(
    self,
    K_features: torch.Tensor,
    V: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
):
    B, H, N, feat_dim = K_features.shape
    chunk_sz = max(1, int(self.landmark_chunk_size))
    k_per_chunk = max(1, int(chunk_sz * self.landmark_ratio))

    with torch.no_grad():
        scores = self.landmark_scorer(K_features).squeeze(-1)  # (B, H, N)

        # Anchor boost for position 0 (always select the first token)
        anchor_boost = torch.zeros_like(scores)
        anchor_boost[:, :, 0] = 1e4
        scores = scores + anchor_boost

        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask.unsqueeze(1), -1e4)

        # Pad to multiple of chunk_size
        pad_len = (chunk_sz - N % chunk_sz) % chunk_sz
        if pad_len > 0:
            scores_pad = F.pad(scores, (0, pad_len), value=-1e4)
        else:
            scores_pad = scores

        num_chunks = scores_pad.shape[-1] // chunk_sz
        scores_c = scores_pad.view(B, H, num_chunks, chunk_sz)

        k_act = min(k_per_chunk, chunk_sz)
        _, local_idx = scores_c.topk(k_act, dim=-1)

        offsets = torch.arange(0, num_chunks * chunk_sz, chunk_sz, device=scores.device).view(
            1, 1, -1, 1
        )
        global_idx = (local_idx + offsets).view(B, H, -1)
        global_idx = global_idx.clamp(max=N - 1)

    K_land = torch.gather(
        K_features,
        2,
        global_idx.unsqueeze(-1).expand(-1, -1, -1, feat_dim),
    )
    V_land = torch.gather(
        V,
        2,
        global_idx.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]),
    )
    return K_land, V_land, global_idx

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

Expand source code

def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
    new_x_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
    x = x.view(new_x_shape)
    return x.permute(0, 2, 1, 3)

class LEKAv2Attention (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_v2", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("leka_v2_attention", supports_causal=True, supports_gqa=True)
class LEKAv2Attention(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_attention_heads = int(num_heads)
        self.num_heads = self.num_attention_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)

        # Hyper-parameters
        self.landmark_ratio = float(
            getattr(comp_config, "landmark_ratio", 0.5) if comp_config else 0.5
        )
        self.landmark_chunk_size = int(
            getattr(comp_config, "landmark_chunk_size", 64) if comp_config else 64
        )
        self.min_landmarks = int(getattr(comp_config, "min_landmarks", 64) if comp_config else 64)
        self.num_landmarks = int(
            getattr(comp_config, "num_landmarks", self.min_landmarks)
            if comp_config
            else self.min_landmarks
        )
        self.feature_dim = int(getattr(comp_config, "feature_dim", 128) if comp_config else 128)
        self.local_window_size = int(
            getattr(comp_config, "local_window_size", 128) if comp_config else 128
        )
        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.k_proj, self.v_proj = self.query, self.key, self.value

        self.register_buffer(
            "rand_proj1",
            torch.randn(self.head_dim, self.feature_dim) / math.sqrt(self.head_dim),
            persistent=True,
        )
        self.register_buffer(
            "rand_proj2",
            torch.randn(self.head_dim, self.feature_dim // 2) / math.sqrt(self.head_dim),
            persistent=True,
        )

        tot_feat = self.feature_dim * 2 + (self.feature_dim // 2) * 2
        self.landmark_scorer = nn.Sequential(
            nn.LayerNorm(tot_feat),
            nn.Linear(tot_feat, 64),
            nn.GELU(),
            nn.Linear(64, 1),
        )

        self.W_corr = nn.Parameter(
            torch.eye(self.head_dim).unsqueeze(0).repeat(self.num_heads, 1, 1)
        )
        self.local_gate = nn.Parameter(torch.full((self.num_heads,), -2.0))
        self.dropout = nn.Dropout(dropout_p)

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

    # ------------------------------------------------------------------
    # Dimension detectors
    # ------------------------------------------------------------------
    @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", "query", "q_lin"):
                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))

    # ------------------------------------------------------------------
    # Kernel feature map (dual-scale RFF)
    # ------------------------------------------------------------------
    def kernel_feature_map(self, x: torch.Tensor) -> torch.Tensor:
        W1 = self.rand_proj1.to(device=x.device, dtype=x.dtype)
        W2 = self.rand_proj2.to(device=x.device, dtype=x.dtype)

        proj1 = torch.matmul(x, W1)
        f1 = torch.cat([torch.cos(proj1), torch.sin(proj1)], dim=-1) / math.sqrt(self.feature_dim)

        proj2 = torch.matmul(x, W2)
        f2 = torch.cat([torch.cos(proj2), torch.sin(proj2)], dim=-1) / math.sqrt(
            self.feature_dim // 2
        )

        return torch.cat([f1, f2], dim=-1)

    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
        new_x_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    # ------------------------------------------------------------------
    # Chunk-wise Landmark Selection (Solves Causal Starvation)
    # ------------------------------------------------------------------
    def select_landmarks_chunked(
        self,
        K_features: torch.Tensor,
        V: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
    ):
        B, H, N, feat_dim = K_features.shape
        chunk_sz = max(1, int(self.landmark_chunk_size))
        k_per_chunk = max(1, int(chunk_sz * self.landmark_ratio))

        with torch.no_grad():
            scores = self.landmark_scorer(K_features).squeeze(-1)  # (B, H, N)

            # Anchor boost for position 0 (always select the first token)
            anchor_boost = torch.zeros_like(scores)
            anchor_boost[:, :, 0] = 1e4
            scores = scores + anchor_boost

            if valid_mask is not None:
                scores = scores.masked_fill(~valid_mask.unsqueeze(1), -1e4)

            # Pad to multiple of chunk_size
            pad_len = (chunk_sz - N % chunk_sz) % chunk_sz
            if pad_len > 0:
                scores_pad = F.pad(scores, (0, pad_len), value=-1e4)
            else:
                scores_pad = scores

            num_chunks = scores_pad.shape[-1] // chunk_sz
            scores_c = scores_pad.view(B, H, num_chunks, chunk_sz)

            k_act = min(k_per_chunk, chunk_sz)
            _, local_idx = scores_c.topk(k_act, dim=-1)

            offsets = torch.arange(0, num_chunks * chunk_sz, chunk_sz, device=scores.device).view(
                1, 1, -1, 1
            )
            global_idx = (local_idx + offsets).view(B, H, -1)
            global_idx = global_idx.clamp(max=N - 1)

        K_land = torch.gather(
            K_features,
            2,
            global_idx.unsqueeze(-1).expand(-1, -1, -1, feat_dim),
        )
        V_land = torch.gather(
            V,
            2,
            global_idx.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]),
        )
        return K_land, V_land, global_idx

    # ------------------------------------------------------------------
    # Block-wise Sliding Window Attention (True O(N) memory)
    # ------------------------------------------------------------------
    def block_local_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, N, d = q.shape
        W = max(1, int(self.local_window_size))

        pad_len = (W - N % W) % W
        if pad_len > 0:
            q_p = F.pad(q, (0, 0, 0, pad_len))
            k_p = F.pad(k, (0, 0, 0, pad_len))
            v_p = F.pad(v, (0, 0, 0, pad_len))
        else:
            q_p, k_p, v_p = q, k, v

        N_pad = q_p.shape[2]
        num_blocks = N_pad // W
        scale = 1.0 / math.sqrt(d)

        q_b = q_p.view(B, H, num_blocks, W, d)
        k_b = k_p.view(B, H, num_blocks, W, d)
        v_b = v_p.view(B, H, num_blocks, W, d)

        # 1. Intra-block (self) attention
        attn_intra = torch.matmul(q_b, k_b.transpose(-1, -2)) * scale
        if self.is_decoder:
            causal_mask = torch.tril(torch.ones(W, W, device=q.device, dtype=torch.bool))
            attn_intra = attn_intra.masked_fill(~causal_mask.view(1, 1, 1, W, W), -1e4)

        # 2. Inter-block (previous block) attention
        k_b_shifted = F.pad(k_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
        v_b_shifted = F.pad(v_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
        attn_inter = torch.matmul(q_b, k_b_shifted.transpose(-1, -2)) * scale
        attn_inter[:, :, 0, :, :] = -1e4  # First block has no previous block

        attn_combined = torch.cat([attn_inter, attn_intra], dim=-1)  # (B, H, nb, W, 2W)

        if valid_mask is not None:
            v_mask_p = F.pad(valid_mask.float(), (0, pad_len), value=0.0).bool()
            # valid_mask shape: (B, N_pad) → broadcast against 2W dim
            v_mask_2w = torch.cat([v_mask_p, v_mask_p], dim=-1).view(B, num_blocks, 1, 2 * W)
            attn_combined = attn_combined.masked_fill(~v_mask_2w.unsqueeze(1), -1e4)

        attn_combined = torch.nan_to_num(attn_combined, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(attn_combined, dim=-1).to(v.dtype)
        attn_weights = self.dropout(attn_weights)

        v_combined = torch.cat([v_b_shifted, v_b], dim=-2)  # (B, H, nb, 2W, d)
        out_local = torch.matmul(attn_weights, v_combined)  # (B, H, nb, W, d)

        return out_local.reshape(B, H, N_pad, d)[:, :, :N, :]

    # ------------------------------------------------------------------
    # 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("LEKAv2._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_()

    # ------------------------------------------------------------------
    # copy_weights (GPT-2 / BERT-nested / DistilBERT / Qwen)
    # ------------------------------------------------------------------
    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, LLaMA-style)
        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"LEKAv2Attention.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(
                            "LEKAv2Attention: 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()

        # Forward sanity check
        try:
            with torch.no_grad():
                d = next(self.parameters()).device
                dt = self.query.down.weight.dtype
                dummy = torch.randn(1, 8, self.dim, device=d, dtype=dt)
                _ = self.forward(dummy)
        except Exception as e:
            logger.debug("LEKAv2 forward sanity check warning: %s", e)

    # ------------------------------------------------------------------
    # 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 = self.num_heads
        device = hidden_states.device

        q = self.query(hidden_states).view(B, N, H, self.head_dim).transpose(1, 2).contiguous()
        k = (
            self.key(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.value(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # RoPE with past_len offset
        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)

        Q_features = self.kernel_feature_map(q)
        K_features = self.kernel_feature_map(k)

        token_valid = None
        if attention_mask is not None:
            token_valid = get_token_valid_mask(attention_mask, B, N, device)

        # 1. Landmark branch (chunk-wise selection, solves causal starvation)
        K_land, V_land, land_idx = self.select_landmarks_chunked(
            K_features,
            v,
            valid_mask=token_valid,
        )
        V_land_adj = torch.einsum("bhld,hde->bhle", V_land, self.W_corr.to(V_land.dtype))

        attn_land = torch.matmul(Q_features.float(), K_land.transpose(-2, -1).contiguous().float())

        if token_valid is not None:
            attn_land = attn_land.masked_fill(~token_valid.unsqueeze(1).unsqueeze(-1), -1e4)

        if self.is_decoder:
            q_pos = torch.arange(N, device=device).view(1, 1, N, 1)
            causal_land_mask = land_idx.unsqueeze(2) <= q_pos
            attn_land = attn_land.masked_fill(~causal_land_mask, -1e4)

        attn_land = softpick(attn_land, dim=-1)
        attn_land = self.dropout(attn_land)
        context_leka = torch.matmul(attn_land.to(v.dtype), V_land_adj)

        # 2. Local branch (block-wise sliding window, true O(N) memory)
        context_local = self.block_local_attention(q, k, v, token_valid)

        # 3. Gate fusion (local dominates at init)
        gate_val = torch.sigmoid(self.local_gate).view(1, H, 1, 1)
        context = (1.0 - gate_val) * context_local + gate_val * context_leka

        # 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)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return out, None

    def get_ortho_loss(self) -> float:
        return 0.0

Every 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 block_local_attention(self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
valid_mask: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def block_local_attention(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    B, H, N, d = q.shape
    W = max(1, int(self.local_window_size))

    pad_len = (W - N % W) % W
    if pad_len > 0:
        q_p = F.pad(q, (0, 0, 0, pad_len))
        k_p = F.pad(k, (0, 0, 0, pad_len))
        v_p = F.pad(v, (0, 0, 0, pad_len))
    else:
        q_p, k_p, v_p = q, k, v

    N_pad = q_p.shape[2]
    num_blocks = N_pad // W
    scale = 1.0 / math.sqrt(d)

    q_b = q_p.view(B, H, num_blocks, W, d)
    k_b = k_p.view(B, H, num_blocks, W, d)
    v_b = v_p.view(B, H, num_blocks, W, d)

    # 1. Intra-block (self) attention
    attn_intra = torch.matmul(q_b, k_b.transpose(-1, -2)) * scale
    if self.is_decoder:
        causal_mask = torch.tril(torch.ones(W, W, device=q.device, dtype=torch.bool))
        attn_intra = attn_intra.masked_fill(~causal_mask.view(1, 1, 1, W, W), -1e4)

    # 2. Inter-block (previous block) attention
    k_b_shifted = F.pad(k_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
    v_b_shifted = F.pad(v_b[:, :, :-1, :, :], (0, 0, 0, 0, 1, 0))
    attn_inter = torch.matmul(q_b, k_b_shifted.transpose(-1, -2)) * scale
    attn_inter[:, :, 0, :, :] = -1e4  # First block has no previous block

    attn_combined = torch.cat([attn_inter, attn_intra], dim=-1)  # (B, H, nb, W, 2W)

    if valid_mask is not None:
        v_mask_p = F.pad(valid_mask.float(), (0, pad_len), value=0.0).bool()
        # valid_mask shape: (B, N_pad) → broadcast against 2W dim
        v_mask_2w = torch.cat([v_mask_p, v_mask_p], dim=-1).view(B, num_blocks, 1, 2 * W)
        attn_combined = attn_combined.masked_fill(~v_mask_2w.unsqueeze(1), -1e4)

    attn_combined = torch.nan_to_num(attn_combined, nan=-1e4, posinf=1e4, neginf=-1e4)
    attn_weights = F.softmax(attn_combined, dim=-1).to(v.dtype)
    attn_weights = self.dropout(attn_weights)

    v_combined = torch.cat([v_b_shifted, v_b], dim=-2)  # (B, H, nb, 2W, d)
    out_local = torch.matmul(attn_weights, v_combined)  # (B, H, nb, W, d)

    return out_local.reshape(B, H, N_pad, d)[:, :, :N, :]

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:
    W1 = self.rand_proj1.to(device=x.device, dtype=x.dtype)
    W2 = self.rand_proj2.to(device=x.device, dtype=x.dtype)

    proj1 = torch.matmul(x, W1)
    f1 = torch.cat([torch.cos(proj1), torch.sin(proj1)], dim=-1) / math.sqrt(self.feature_dim)

    proj2 = torch.matmul(x, W2)
    f2 = torch.cat([torch.cos(proj2), torch.sin(proj2)], dim=-1) / math.sqrt(
        self.feature_dim // 2
    )

    return torch.cat([f1, f2], dim=-1)

def select_landmarks_chunked(self,
K_features: torch.Tensor,
V: torch.Tensor,
valid_mask: torch.Tensor | None = None)

Expand source code

def select_landmarks_chunked(
    self,
    K_features: torch.Tensor,
    V: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
):
    B, H, N, feat_dim = K_features.shape
    chunk_sz = max(1, int(self.landmark_chunk_size))
    k_per_chunk = max(1, int(chunk_sz * self.landmark_ratio))

    with torch.no_grad():
        scores = self.landmark_scorer(K_features).squeeze(-1)  # (B, H, N)

        # Anchor boost for position 0 (always select the first token)
        anchor_boost = torch.zeros_like(scores)
        anchor_boost[:, :, 0] = 1e4
        scores = scores + anchor_boost

        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask.unsqueeze(1), -1e4)

        # Pad to multiple of chunk_size
        pad_len = (chunk_sz - N % chunk_sz) % chunk_sz
        if pad_len > 0:
            scores_pad = F.pad(scores, (0, pad_len), value=-1e4)
        else:
            scores_pad = scores

        num_chunks = scores_pad.shape[-1] // chunk_sz
        scores_c = scores_pad.view(B, H, num_chunks, chunk_sz)

        k_act = min(k_per_chunk, chunk_sz)
        _, local_idx = scores_c.topk(k_act, dim=-1)

        offsets = torch.arange(0, num_chunks * chunk_sz, chunk_sz, device=scores.device).view(
            1, 1, -1, 1
        )
        global_idx = (local_idx + offsets).view(B, H, -1)
        global_idx = global_idx.clamp(max=N - 1)

    K_land = torch.gather(
        K_features,
        2,
        global_idx.unsqueeze(-1).expand(-1, -1, -1, feat_dim),
    )
    V_land = torch.gather(
        V,
        2,
        global_idx.unsqueeze(-1).expand(-1, -1, -1, V.shape[-1]),
    )
    return K_land, V_land, global_idx

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

Expand source code

def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
    new_x_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
    x = x.view(new_x_shape)
    return x.permute(0, 2, 1, 3)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class LSHAttention (config: Any,
rank: int | None = None,
rounds: int = 4,
num_buckets: int = 16,
lora_rank: int = 16,
lora_alpha: float = 32.0,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
is_decoder: bool | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("lsh", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("lsh_attention", supports_causal=True, supports_gqa=True)
class LSHAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any,
        rank: int | None = None,
        rounds: int = 4,
        num_buckets: int = 16,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # LSH-02: robust super().__init__
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        # LSH-10: set self.config explicitly
        self.config = config

        # LSH-03, LSH-13: robust dimension detection (independent of ModelAdapter)
        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.hidden_size = int(hidden_size)
        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.kv_head_size = self.num_kv_heads * self.head_dim
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))

        # LSH-11: clamp num_buckets
        self.rounds = max(1, int(getattr(config, "lsh_rounds", rounds) if config else rounds))
        self.num_buckets = max(
            2,
            int(getattr(config, "lsh_buckets", num_buckets) if config else num_buckets),
        )

        # Decoder detection
        self.is_decoder = _detect_decoder(config, is_decoder, source_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # LoRA config
        self.lora_rank = int(
            (getattr(config, "lora_rank", None) if config else None)
            or (getattr(comp_config, "lora_rank", None) if comp_config else None)
            or lora_rank
        )
        self.lora_alpha = float(
            (getattr(config, "lora_alpha", None) if config else None)
            or (getattr(comp_config, "lora_alpha", None) if comp_config else None)
            or lora_alpha
        )

        # LSH-08: projection rank
        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)
        )

        # FactorizedLoRALinear projections
        FF = FactorizedLoRALinear
        self.query = FF(
            self.hidden_size,
            self.all_head_size,
            rank=self.proj_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )
        self.key = FF(
            self.hidden_size,
            self.kv_head_size,
            rank=self.proj_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )
        self.value = FF(
            self.hidden_size,
            self.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.hidden_size,
            rank=self.proj_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )

        dropout_p = float(
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                ),
            )
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        # LSH-14: per-round random hyperplane projection, normalized per column
        # Shape: (rounds, head_dim, num_buckets // 2)
        half_buckets = max(1, self.num_buckets // 2)
        proj_data = torch.randn(self.rounds, self.head_dim, half_buckets)
        # Normalize each hyperplane (column) to unit L2 norm
        proj_data = F.normalize(proj_data, p=2, dim=1)
        self.register_buffer("proj_matrices", proj_data, persistent=False)

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

    # ------------------------------------------------------------------
    # LSH-13: 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))

    # ------------------------------------------------------------------
    # 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("LSHAttention._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_()

    # ------------------------------------------------------------------
    # LSH-07: 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),
        }

        # Try UniversalAdapter first
        try:
            from neuralops.core.adapters import UniversalAdapter

            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
        except Exception:
            pass

        # 1. GPT-2 fused c_attn
        if (
            hasattr(src, "c_attn")
            and src.c_attn is not None
            and any(extracted[k][0] is None for k in ("q", "k", "v"))
        ):
            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 = 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, :],
                    c_b[:split] if c_b is not None else None,
                )
                extracted["k"] = (
                    c_w[split : 2 * split, :],
                    c_b[split : 2 * split] if c_b is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :, :],
                    c_b[2 * split :] if c_b is not None else None,
                )

            if hasattr(src, "c_proj") and src.c_proj is not None and extracted["o"][0] is 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")
            and any(extracted[k][0] is None for k in ("q", "k", "v"))
        ):
            s = src.self
            if hasattr(s, "query") and extracted["q"][0] is None:
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key") and extracted["k"][0] is None:
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value") and extracted["v"][0] is None:
                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))

        # 3. Flat modules (DistilBERT, Qwen, LLaMA)
        else:
            named_map = {
                "q": ["q_proj", "q_lin", "query"],
                "k": ["k_proj", "k_lin", "key"],
                "v": ["v_proj", "v_lin", "value"],
                "o": ["o_proj", "out_proj", "out_lin", "dense"],
            }
            for proj, names in named_map.items():
                if extracted[proj][0] is None:
                    for name in names:
                        if hasattr(src, name):
                            mod = getattr(src, name)
                            if mod is not None and hasattr(mod, "weight"):
                                w_data = mod.weight.data
                                if (
                                    proj == "o"
                                    and w_data.ndim == 2
                                    and w_data.shape[0] != self.hidden_size
                                    and w_data.shape[1] == self.hidden_size
                                ):
                                    w_data = w_data.t().contiguous()
                                extracted[proj] = (w_data, getattr(mod, "bias", 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:
            raise RuntimeError(
                f"LSHAttention.copy_weights: missing mandatory projections {missing} "
                f"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.linalg.norm(mod.weight.float()).item())
                    if w_norm < 1e-4:
                        logger.warning(
                            "LSHAttention: 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()

        # Forward sanity check
        try:
            with torch.no_grad():
                d = next(self.parameters()).device
                dt = self.query.down.weight.dtype
                dummy = torch.randn(1, 8, self.hidden_size, device=d, dtype=dt)
                _ = self.forward(dummy)
        except Exception as e:
            logger.debug("LSHAttention forward sanity check warning: %s", e)

    # ------------------------------------------------------------------
    # LSH hashing
    # ------------------------------------------------------------------
    def _hash(self, vectors: torch.Tensor, proj: torch.Tensor) -> torch.Tensor:
        """
        Hash vectors via random hyperplane sign pattern.
        vectors: (B, H, N, d); proj: (d, num_buckets//2)
        Returns: (B, H, N) bucket indices in [0, num_buckets)
        """
        norm = F.normalize(vectors.float(), p=2, dim=-1)
        dots = torch.matmul(norm, proj.float())  # (B, H, N, num_buckets//2)
        signs = (dots >= 0).long()  # (B, H, N, num_buckets//2)
        # Combine sign bits into a single bucket index
        num_planes = signs.shape[-1]
        weights = (2 ** torch.arange(num_planes, device=signs.device)).long()
        bucket = (signs * weights).sum(dim=-1)  # (B, H, N)
        bucket = bucket % self.num_buckets
        return bucket

    # ------------------------------------------------------------------
    # LSH-09: chunked LSH attention
    # ------------------------------------------------------------------
    def _lsh_attention_chunked(
        self,
        Q: torch.Tensor,  # (B, H, N, d)
        K: torch.Tensor,  # (B, H, N, d)
        V: torch.Tensor,  # (B, H, N, d)
        valid_mask: torch.Tensor | None,  # (B, N) bool
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        Chunked LSH attention. For each round, compute per-query chunk attention
        over keys in the same bucket. Memory O(B*H*N*C) per round (C=chunk).
        """
        B, H, N, d = Q.shape
        device = Q.device
        dtype = Q.dtype
        scale = d**-0.5

        num_chunks = max(1, (N + 63) // 64)
        chunk_size = (N + num_chunks - 1) // num_chunks

        accum = torch.zeros_like(Q)
        rounds_eff = min(self.rounds, self.proj_matrices.size(0))

        for r in range(rounds_eff):
            proj_r = self.proj_matrices[r].to(device=device, dtype=dtype)
            q_hash = self._hash(Q, proj_r)  # (B, H, N)
            k_hash = self._hash(K, proj_r)

            for c in range(num_chunks):
                c_start = c * chunk_size
                c_end = min(c_start + chunk_size, N)
                Qc = Q[:, :, c_start:c_end]  # (B, H, C, d)
                hq_c = q_hash[:, :, c_start:c_end]  # (B, H, C)

                # bucket match: (B, H, C, N)
                bucket_match = hq_c.unsqueeze(-1) == k_hash.unsqueeze(-2)

                # Mask padding keys
                if valid_mask is not None:
                    kv_mask = valid_mask[:, None, None, :]  # (B, 1, 1, N)
                    bucket_match = bucket_match & kv_mask

                # scores: (B, H, C, N)
                scores_c = (
                    torch.matmul(Qc.float(), K.float().transpose(-2, -1).contiguous()) * scale
                )

                scores_c = scores_c.masked_fill(~bucket_match, -1e4)
                scores_c = torch.nan_to_num(scores_c, nan=-1e4, posinf=1e4, neginf=-1e4)

                attn_c = F.softmax(scores_c, dim=-1).to(dtype)
                attn_c = torch.nan_to_num(attn_c, nan=0.0)
                attn_c = self.dropout(attn_c)

                accum[:, :, c_start:c_end] = accum[:, :, c_start:c_end] + torch.matmul(attn_c, V)

        accum = accum / float(rounds_eff)
        return accum, None

    # ------------------------------------------------------------------
    # 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]:
        # Handle key_value_states via kwargs
        key_value_states = kwargs.get("key_value_states", encoder_hidden_states)
        is_cross = key_value_states is not None
        kv_input = key_value_states if is_cross else hidden_states

        B, N_q, _ = 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

        Q = self.query(hidden_states).view(B, N_q, H, d).transpose(1, 2).contiguous()
        K = self.key(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V = self.value(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # 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 = 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_q, :]
                sin = sin[..., past_len : past_len + N_q, :]
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # GQA
        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        # Valid mask
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        valid_mask = None
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    valid_mask = eff_mask
                elif eff_mask.dtype.is_floating_point:
                    valid_mask = eff_mask > -1e4
                else:
                    valid_mask = eff_mask > 0
            elif eff_mask.dim() == 4:
                if eff_mask.dtype == torch.bool:
                    v = eff_mask
                elif eff_mask.dtype.is_floating_point:
                    v = eff_mask > -1e4
                else:
                    v = eff_mask > 0
                valid_mask = v.any(dim=-2).squeeze(1)

        # LSH-15: decoders fall back to exact attention (LSH doesn't naturally support causal)
        use_exact = self.is_decoder or is_cross or N_q <= self.num_buckets

        if use_exact:
            scale = d**-0.5
            scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

            if valid_mask is not None:
                scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

            if self.is_decoder and not is_cross:
                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, -1e4)

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

            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
                attn = attn * hm.to(attn.dtype)

            context = torch.matmul(attn, V)
            attn_ret = attn if output_attentions else None
        else:
            # LSH-09: chunked LSH attention
            context, _ = self._lsh_attention_chunked(Q, K, V, valid_mask)
            attn_ret = None
            # Note: LSH path does not produce a single interpretable attention matrix.

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

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

        return out.to(dtype), attn_ret

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class LinformerAttention (config: Any | None = None,
rank: int | None = None,
k: int | None = None,
max_seq_len: int | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("linformer", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("linformer_attention", supports_causal=True, supports_gqa=True)
class LinformerAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "linear_attn"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        k: int | None = None,
        max_seq_len: int | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # LF-02: robust super().__init__
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        # LF-04: robust dimension detection
        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

        # Linformer hyperparams
        k_val = (
            k
            if k is not None
            else getattr(
                comp_config,
                "linformer_k",
                getattr(comp_config, "k", None) if comp_config else None,
            )
        )
        if k_val is None:
            k_val = getattr(config, "linformer_k", getattr(config, "k", 64)) if config else 64
        self.k_proj_dim = int(k_val)

        pos_emb = getattr(config, "max_position_embeddings", None) if config else None
        if pos_emb is None:
            pos_emb = max_seq_len if max_seq_len is not None else 512
        self.max_seq_len = max(512, int(pos_emb))

        # LF-11: centralized decoder detection
        self.is_decoder = _detect_decoder(
            config, is_decoder or kwargs.get("is_causal"), source_attn
        )
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # LoRA config
        self.lora_rank = int(
            (getattr(config, "lora_rank", None) if config else 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) if config else None)
            or (getattr(comp_config, "lora_alpha", None) if comp_config else None)
            or kwargs.get("lora_alpha", 32.0)
        )

        # LF-12: projection rank for FactorizedLoRALinear
        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)
        )

        # FactorizedLoRALinear projections
        FF = FactorizedLoRALinear
        self.q_proj = FF(
            self.dim,
            self.all_head_size,
            rank=self.proj_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )
        self.k_proj = 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.v_proj = 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,
        )

        # Linformer projections: E, F: (k, max_seq_len)
        self.E = nn.Parameter(torch.empty(self.k_proj_dim, self.max_seq_len))
        self.F = nn.Parameter(torch.empty(self.k_proj_dim, self.max_seq_len))
        with torch.no_grad():
            nn.init.normal_(self.E, std=0.02)
            nn.init.normal_(self.F, std=0.02)
            # LF-15: only touch the min(k, max_seq) block
            min_dim = min(self.k_proj_dim, self.max_seq_len)
            self.E.data[:min_dim, :min_dim].fill_diagonal_(1.0)
            self.F.data[:min_dim, :min_dim].fill_diagonal_(1.0)

        dropout_p = float(
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                ),
            )
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)
        self.triton_kwargs = {"phi": "identity", "eps": 1e-6}

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

    # ------------------------------------------------------------------
    # LF-04: 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))

    # ------------------------------------------------------------------
    # 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 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

    # ------------------------------------------------------------------
    # SVD init helper
    # ------------------------------------------------------------------
    @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("Linformer._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_()

    # ------------------------------------------------------------------
    # copy_weights (LF-07, LF-08, LF-11)
    # ------------------------------------------------------------------
    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, Qwen, LLaMA)
        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"LinformerAttention.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):
            """LF-07: adapt w for GQA: average down or repeat to match 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.q_proj)
            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.k_proj.out_features)
                self._svd_init(w_k_res, None, self.k_proj)
            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.v_proj.out_features)
                self._svd_init(w_v_res, None, self.v_proj)
            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)

            # LF-08: norm verification via float() cast
            for name, mod in (
                ("q_proj", self.q_proj),
                ("k_proj", self.k_proj),
                ("v_proj", self.v_proj),
                ("out_proj", self.out_proj),
            ):
                try:
                    w_norm = float(torch.linalg.norm(mod.weight.float()).item())
                    if w_norm < 1e-4:
                        logger.warning(
                            "LinformerAttention: weight '%s' has near-zero norm (%.2e).",
                            name,
                            w_norm,
                        )
                except Exception:
                    pass

            for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
                mod.freeze_base_weights()

        # LF-11: inherit is_decoder
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)
            self.is_causal = self.is_decoder

        if hasattr(src, "rotary_emb") and src.rotary_emb is not None:
            self.rotary_emb = src.rotary_emb
        else:
            if not hasattr(self, "rotary_emb"):
                self.rotary_emb = None

        # LF-13: forward sanity check
        try:
            with torch.no_grad():
                d = next(self.parameters()).device
                dt = self.q_proj.down.weight.dtype
                dummy = torch.randn(1, 8, self.dim, device=d, dtype=dt)
                _ = self.forward(dummy)
        except Exception as e:
            logger.debug("LinformerAttention forward sanity check warning: %s", e)

    # ------------------------------------------------------------------
    # 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
        device = hidden_states.device
        dtype = hidden_states.dtype

        q = self.q_proj(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous()
        k = self.k_proj(hidden_states).view(B, N, self.num_kv_heads, d).transpose(1, 2).contiguous()
        v = self.v_proj(hidden_states).view(B, N, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # LF-05: RoPE with past_len offset
        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
            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)

        # Mask K, V (LF-14: shape-safe)
        valid_mask = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                if attention_mask.dtype == torch.bool:
                    valid_mask = attention_mask
                elif attention_mask.dtype.is_floating_point:
                    valid_mask = attention_mask > -1e4
                else:
                    valid_mask = attention_mask > 0
            elif attention_mask.dim() == 4:
                if attention_mask.dtype == torch.bool:
                    v_tmp = attention_mask
                elif attention_mask.dtype.is_floating_point:
                    v_tmp = attention_mask > -1e4
                else:
                    v_tmp = attention_mask > 0
                valid_mask = v_tmp.any(dim=-2).squeeze(1)
            else:
                valid_mask = torch.ones(B, N, device=device, dtype=torch.bool)

        if valid_mask is not None:
            mask_expanded = valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
            k = k * mask_expanded
            v = v * mask_expanded
        else:
            mask_expanded = torch.ones(B, 1, N, 1, device=device, dtype=dtype)

        L = N
        k_eff = min(self.k_proj_dim, L)

        # E, F: (k_proj_dim, max_seq_len) → (k_eff, L)
        if self.max_seq_len >= L:
            E_act = self.E[:k_eff, :L].to(dtype=dtype, device=device)
            F_act = self.F[:k_eff, :L].to(dtype=dtype, device=device)
        else:
            # LF-10: interpolation with 3D input
            E_act = (
                F.interpolate(
                    self.E[:k_eff].unsqueeze(0).float(),
                    size=L,
                    mode="linear",
                    align_corners=False,
                )
                .squeeze(0)
                .to(dtype=dtype, device=device)
            )
            F_act = (
                F.interpolate(
                    self.F[:k_eff].unsqueeze(0).float(),
                    size=L,
                    mode="linear",
                    align_corners=False,
                )
                .squeeze(0)
                .to(dtype=dtype, device=device)
            )

        # Project key/value sequence: (B, H, k_eff, d)
        K_proj = torch.matmul(E_act.to(k.dtype), k)
        V_proj = torch.matmul(F_act.to(v.dtype), v)

        scale = 1.0 / math.sqrt(d)
        scores = torch.matmul(q.float(), K_proj.transpose(-2, -1).contiguous().float()) * scale

        # LF-09: proper causal masking for projected K/V
        if self.is_decoder:
            if self.k_proj_dim >= L:
                # Exact causal on full sequence — recompute without projection
                causal_scores = (
                    torch.matmul(q.float(), k.transpose(-2, -1).contiguous().float()) * scale
                )
                causal_mask = torch.tril(torch.ones(L, L, device=device, dtype=torch.bool))
                scores = causal_scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4)
                V_proj = v
            else:
                # Approximate causal on projected K/V: mask each projected key by its mean position
                torch.tril(torch.ones(L, k_eff, device=device, dtype=torch.bool))
                # For each query position i, allow projected keys whose support range includes <= i
                # Simplified: allow key j if j <= i (proportional to support)
                idx_q = torch.arange(L, device=device).unsqueeze(1)  # (L, 1)
                idx_k = torch.arange(k_eff, device=device).unsqueeze(0)  # (1, k_eff)
                # Each projected key j covers roughly L/k_eff positions
                step = max(1, L // k_eff)
                allowed = (idx_k * step) <= idx_q  # (L, k_eff)
                scores = scores.masked_fill(~allowed.unsqueeze(0).unsqueeze(0), -1e4)

        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask.unsqueeze(1).unsqueeze(-1), -1e4)

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

        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
            attn = attn * hm.to(attn.dtype)

        context = torch.matmul(attn, V_proj.to(attn.dtype))
        if valid_mask is not None:
            context = context * mask_expanded
        context = context.transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return out.to(dtype), (attn if output_attentions else None)

Every 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.

var triton_kernel : str

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

class LoLA (config: Any | None = None,
model_config: Any | None = None,
rank: int | None = None,
lora_rank: int = 16,
dropout: float = 0.1,
comp_config: Any | None = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("lola_attn", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("lola_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("lola", supports_causal=True, supports_gqa=True)
class LoLAAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        rank: int | None = None,
        lora_rank: int = 16,
        dropout: float = 0.1,
        comp_config: Any | None = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config if config is not None else model_config

        def _cfg_get(name: str, fallback: Any) -> Any:
            if cfg is not None:
                val = getattr(cfg, name, None)
                if val is not None:
                    return val
            return fallback

        actual_rank = rank if rank is not None else _cfg_get("rank", 32)
        super().__init__(cfg, actual_rank)

        cfg_head_dim = getattr(cfg, "head_dim", None) if cfg is not None else None
        if cfg is not None:
            self.adapter = ModelAdapter(cfg)
            self.num_heads = self.adapter.num_heads or getattr(cfg, "num_attention_heads", 12)
            self.num_kv_heads = self.adapter.num_key_value_heads or getattr(
                cfg, "num_key_value_heads", self.num_heads
            )
            self.hidden_size = self.adapter.hidden_size or getattr(cfg, "hidden_size", 768)
            self.head_dim = int(
                cfg_head_dim
                or kwargs.get("head_dim")
                or self.adapter.head_dim
                or (self.hidden_size // self.num_heads)
            )
            self.kv_repeat = self.adapter.kv_repeat or 1
        else:
            self.adapter = None
            self.num_heads = kwargs.get("num_heads", kwargs.get("num_attention_heads", 12))
            self.num_kv_heads = kwargs.get("num_key_value_heads", self.num_heads)
            self.hidden_size = kwargs.get("hidden_size", 768)
            self.head_dim = int(kwargs.get("head_dim") or (self.hidden_size // self.num_heads))
            self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))

        _mt = str(getattr(cfg, "model_type", "") or "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        self.is_decoder = (
            (
                getattr(cfg, "is_decoder", False)
                or getattr(cfg, "causal", False)
                or any(_mt.startswith(p) for p in _decoder_prefixes)
                or bool(kwargs.get("is_decoder", False))
            )
            if cfg
            else bool(kwargs.get("is_decoder", False))
        )

        self.all_head_size: int = self.num_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.is_nested_encoder: bool = False

        self.rank: int = actual_rank
        self.lora_rank: int = _cfg_get("lora_rank", lora_rank)

        self.q_proj = LowRankLinear(
            self.hidden_size,
            self.all_head_size,
            self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.k_proj = LowRankLinear(
            self.hidden_size,
            self.kv_head_size,
            self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.v_proj = LowRankLinear(
            self.hidden_size,
            self.kv_head_size,
            self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.out_proj = LowRankLinear(
            self.all_head_size,
            self.hidden_size,
            self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )

        self.dropout: nn.Dropout = nn.Dropout(dropout)

        max_seq_len: int = (
            getattr(comp_config, "max_seq_len", 128) if comp_config is not None else 128
        )
        low_rank_cache_size: int = (
            getattr(comp_config, "low_rank_cache_size", 16) if comp_config is not None else 16
        )
        self.cache: LowRankCache = LowRankCache(max_len=max_seq_len, low_rank_k=low_rank_cache_size)
        self.LR_scale: nn.Parameter = nn.Parameter(torch.tensor(4.0))
        self.sparse_gate: nn.Parameter = nn.Parameter(torch.tensor(-3.0))
        self.rotary_emb = (
            getattr(pretrained_attn, "rotary_emb", None) if pretrained_attn is not None else None
        )

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

    def reset_cache(self) -> None:
        if hasattr(self, "cache") and hasattr(self.cache, "reset"):
            self.cache.reset()

    def _set_lowrank_factors(
        self,
        lowrank: LowRankLinear,
        weight: torch.Tensor,
        bias: torch.Tensor | None,
    ) -> None:
        with torch.no_grad():
            W = weight.float()
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            r = min(self.rank, U.size(1), Vt.size(0))
            sqrt_s = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            lowrank.U.data.zero_()
            lowrank.V.data.zero_()
            out_d, in_d = lowrank.U.size(0), lowrank.V.size(1)
            lowrank.U.data[: min(out_d, U.size(0)), :r].copy_(
                (U[: min(out_d, U.size(0)), :r] * sqrt_s).to(lowrank.U.dtype)
            )
            lowrank.V.data[:r, : min(in_d, Vt.size(1))].copy_(
                (sqrt_s.unsqueeze(1) * Vt[:r, : min(in_d, Vt.size(1))]).to(lowrank.V.dtype)
            )
            if lowrank.bias is not None and bias is not None:
                b_val = bias.data if isinstance(bias, nn.Parameter) else bias
                lowrank.bias.data[: min(len(lowrank.bias), len(b_val))].copy_(
                    b_val[: min(len(lowrank.bias), len(b_val))]
                )

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

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        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

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            c_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (weight.dim() == 2 and weight.shape[0] < weight.shape[1])
                ):
                    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:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                p_cls = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__
                b_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or ("Conv1D" in b_cls)
                    or hasattr(c_proj_mod, "nf")
                    or hasattr(base_o, "nf")
                    or (
                        w_o.dim() == 2
                        and (
                            w_o.shape[0] < w_o.shape[1]
                            or (
                                w_o.shape[0] == self.hidden_size
                                and w_o.shape[1] == self.hidden_size
                            )
                        )
                    )
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                extracted["o"] = (w_o, b_o.data.clone() if b_o is not None else None)

        # RoBERTa / BERT nested
        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
            for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]:
                mod = getattr(s, sub_name, None)
                if mod is not None:
                    base_mod = getattr(mod, "base_layer", mod)
                    if hasattr(base_mod, "weight"):
                        extracted[proj_key] = (
                            base_mod.weight.data.clone(),
                            (
                                base_mod.bias.data.clone()
                                if getattr(base_mod, "bias", None) is not None
                                else None
                            ),
                        )

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            dense = src.output.dense
            base_dense = getattr(dense, "base_layer", dense)
            if hasattr(base_dense, "weight"):
                extracted["o"] = (
                    base_dense.weight.data.clone(),
                    (
                        base_dense.bias.data.clone()
                        if getattr(base_dense, "bias", None) is not None
                        else None
                    ),
                )

        # Modern decoders + DistilBERT
        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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)
                        base_l = getattr(layer, "base_layer", layer)
                        if hasattr(base_l, "weight"):
                            extracted[proj] = (
                                base_l.weight.data.clone(),
                                (
                                    base_l.bias.data.clone()
                                    if getattr(base_l, "bias", None) is not None
                                    else None
                                ),
                            )
                            break

        # Nested encoder detection
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            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("Module copy_weights warning. Proceeding with initialized weights.")

        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"LoLAAttention extracted weight for {proj_key} has "
                        f"near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"LoLAAttention extracted weight for output projection has "
                    f"near-zero Frobenius norm ({norm_o.item():.6f})"
                )

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

        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    self._set_lowrank_factors(target, w_oriented, b)
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.U.size(0), target.V.size(1), self.rank)
                    target.U.data.zero_()
                    target.V.data.zero_()
                    target.U.data[:min_d, :min_d].fill_diagonal_(1.0)
                    target.V.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

        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

        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()

        # RoPE
        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"))

        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = d**-0.5

        _cfg_mt = (
            str(getattr(self.config, "model_type", "") or "").lower()
            if self.config is not None
            else ""
        )
        is_decoder = (
            self.is_decoder
            or (getattr(self.config, "is_decoder", False) if self.config is not None else False)
            or any(
                _cfg_mt.startswith(p)
                for p in (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "gemma",
                    "phi",
                    "opt",
                    "bloom",
                    "falcon",
                    "olmo",
                    "mpt",
                    "stablelm",
                )
            )
        )

        scores_lr = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale

        mask = attention_mask if attention_mask is not None else kwargs.get("mask")
        if mask is not None:
            if mask.dim() == 4:
                if mask.dtype == torch.bool:
                    scores_lr = scores_lr.masked_fill(~mask, -10000.0)
                else:
                    scores_lr = scores_lr + mask.to(scores_lr.dtype)
            elif mask.dim() == 3:
                scores_lr = scores_lr + mask.unsqueeze(1).to(scores_lr.dtype)
            elif mask.dim() == 2:
                valid_m = get_token_valid_mask(mask, B, N_k, device)
                if valid_m is not None:
                    scores_lr = scores_lr.masked_fill(
                        (~valid_m).unsqueeze(1).unsqueeze(2), -10000.0
                    )
            else:
                valid_m = get_token_valid_mask(mask, B, N_k, device)
                if valid_m is not None:
                    scores_lr = scores_lr.masked_fill(
                        (~valid_m).unsqueeze(1).unsqueeze(2), -10000.0
                    )

        if is_decoder and N_k >= N_q:
            causal_tril = (
                torch.tril(
                    torch.ones((N_q, N_k), device=device, dtype=torch.bool),
                    diagonal=N_k - N_q,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores_lr = scores_lr.masked_fill(~causal_tril, -10000.0)

        # =============================================================
        # LL-FIX-01/02/03: cache only for streaming decoders with past_key_value.
        # For encoders (bidirectional), the cache would mix tokens across batches
        # and hold large view-retained buffers, causing ~70% VRAM bloat.
        # =============================================================
        use_cache = bool(is_decoder and past_key_value is not None and not is_cross_attention)

        if use_cache:
            self.cache.update(K, V)
            cached_k, cached_v = self.cache.get_cached()
        else:
            # Ensure stale state is dropped so batches stay isolated.
            self.cache.reset()
            cached_k, cached_v = None, None

        if cached_k is not None and cached_v is not None and cached_k.size(0) == (B * H):
            C = cached_k.size(1)
            K_cache = cached_k.view(B, H, C, d)
            V_cache = cached_v.view(B, H, C, d)
            scores_cache = (
                torch.matmul(Q, K_cache.transpose(-2, -1).contiguous())
                * scale
                * torch.sigmoid(torch.clamp(self.sparse_gate, min=-10.0, max=10.0))
            )
            scores = torch.cat(
                [
                    scores_lr * torch.sigmoid(torch.clamp(self.LR_scale, min=-10.0, max=10.0)),
                    scores_cache,
                ],
                dim=-1,
            )
            V_comb = torch.cat([V, V_cache], dim=2)
        else:
            scores = scores_lr
            V_comb = V

        attn = torch.nan_to_num(F.softmax(scores.float(), dim=-1), nan=0.0).to(dtype)
        attn = self.dropout(attn)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn = attn * head_mask

        context = torch.matmul(attn, V_comb)
        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)

        return output, (attn if output_attentions else None)

Every 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.

var supports_rope : bool

The type of the None singleton.

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    from neuralops.core.adapters import UniversalAdapter

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

    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

    # GPT-2 fused c_attn
    if hasattr(src, "c_attn") and (
        extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
    ):
        c_attn_mod = src.c_attn
        base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
        weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
        if weight is not None:
            weight = weight.data.clone()
            bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            bias_data = bias.data.clone() if bias is not None else None
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or (weight.dim() == 2 and weight.shape[0] < weight.shape[1])
            ):
                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:
        c_proj_mod = src.c_proj
        base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
        w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
        if w_o is not None:
            w_o = w_o.data.clone()
            p_cls = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__
            b_cls = getattr(base_o, "__class__", type(base_o)).__name__
            if (
                ("Conv1D" in p_cls)
                or ("Conv1D" in b_cls)
                or hasattr(c_proj_mod, "nf")
                or hasattr(base_o, "nf")
                or (
                    w_o.dim() == 2
                    and (
                        w_o.shape[0] < w_o.shape[1]
                        or (
                            w_o.shape[0] == self.hidden_size
                            and w_o.shape[1] == self.hidden_size
                        )
                    )
                )
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
            extracted["o"] = (w_o, b_o.data.clone() if b_o is not None else None)

    # RoBERTa / BERT nested
    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
        for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]:
            mod = getattr(s, sub_name, None)
            if mod is not None:
                base_mod = getattr(mod, "base_layer", mod)
                if hasattr(base_mod, "weight"):
                    extracted[proj_key] = (
                        base_mod.weight.data.clone(),
                        (
                            base_mod.bias.data.clone()
                            if getattr(base_mod, "bias", None) is not None
                            else None
                        ),
                    )

    if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
        dense = src.output.dense
        base_dense = getattr(dense, "base_layer", dense)
        if hasattr(base_dense, "weight"):
            extracted["o"] = (
                base_dense.weight.data.clone(),
                (
                    base_dense.bias.data.clone()
                    if getattr(base_dense, "bias", None) is not None
                    else None
                ),
            )

    # Modern decoders + DistilBERT
    named_map = {
        "q": ["q_proj", "q_lin", "query"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    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)
                    base_l = getattr(layer, "base_layer", layer)
                    if hasattr(base_l, "weight"):
                        extracted[proj] = (
                            base_l.weight.data.clone(),
                            (
                                base_l.bias.data.clone()
                                if getattr(base_l, "bias", None) is not None
                                else None
                            ),
                        )
                        break

    # Nested encoder detection
    cls_name = src.__class__.__name__
    has_out_proj = any(
        hasattr(src, attr)
        for attr in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    is_nested = False
    if (
        hasattr(src, "self")
        and hasattr(src.self, "query")
        or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
            )
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        )
        and not has_out_proj
    ):
        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("Module copy_weights warning. Proceeding with initialized weights.")

    for proj_key in ["q", "k", "v"]:
        w = extracted[proj_key][0]
        if w is not None:
            norm = torch.linalg.norm(w.float())
            if norm < 1e-4:
                logger.warning(
                    f"LoLAAttention extracted weight for {proj_key} has "
                    f"near-zero Frobenius norm ({norm.item():.6f})"
                )

    if extracted["o"][0] is not None:
        norm_o = torch.linalg.norm(extracted["o"][0].float())
        if norm_o < 1e-4:
            logger.warning(
                f"LoLAAttention extracted weight for output projection has "
                f"near-zero Frobenius norm ({norm_o.item():.6f})"
            )

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

    with torch.no_grad():
        for proj_key, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                w_oriented = _orient(w, in_dim)
                self._set_lowrank_factors(target, w_oriented, b)
            elif proj_key == "o" and self.is_nested_encoder:
                min_d = min(target.U.size(0), target.V.size(1), self.rank)
                target.U.data.zero_()
                target.V.data.zero_()
                target.U.data[:min_d, :min_d].fill_diagonal_(1.0)
                target.V.data[:min_d, :min_d].fill_diagonal_(1.0)
                if target.bias is not None:
                    target.bias.data.zero_()

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

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,
*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,
    *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

    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()

    # RoPE
    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"))

    if self.kv_repeat > 1:
        K = repeat_kv(K, self.kv_repeat)
        V = repeat_kv(V, self.kv_repeat)

    scale = d**-0.5

    _cfg_mt = (
        str(getattr(self.config, "model_type", "") or "").lower()
        if self.config is not None
        else ""
    )
    is_decoder = (
        self.is_decoder
        or (getattr(self.config, "is_decoder", False) if self.config is not None else False)
        or any(
            _cfg_mt.startswith(p)
            for p in (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
        )
    )

    scores_lr = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale

    mask = attention_mask if attention_mask is not None else kwargs.get("mask")
    if mask is not None:
        if mask.dim() == 4:
            if mask.dtype == torch.bool:
                scores_lr = scores_lr.masked_fill(~mask, -10000.0)
            else:
                scores_lr = scores_lr + mask.to(scores_lr.dtype)
        elif mask.dim() == 3:
            scores_lr = scores_lr + mask.unsqueeze(1).to(scores_lr.dtype)
        elif mask.dim() == 2:
            valid_m = get_token_valid_mask(mask, B, N_k, device)
            if valid_m is not None:
                scores_lr = scores_lr.masked_fill(
                    (~valid_m).unsqueeze(1).unsqueeze(2), -10000.0
                )
        else:
            valid_m = get_token_valid_mask(mask, B, N_k, device)
            if valid_m is not None:
                scores_lr = scores_lr.masked_fill(
                    (~valid_m).unsqueeze(1).unsqueeze(2), -10000.0
                )

    if is_decoder and N_k >= N_q:
        causal_tril = (
            torch.tril(
                torch.ones((N_q, N_k), device=device, dtype=torch.bool),
                diagonal=N_k - N_q,
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        scores_lr = scores_lr.masked_fill(~causal_tril, -10000.0)

    # =============================================================
    # LL-FIX-01/02/03: cache only for streaming decoders with past_key_value.
    # For encoders (bidirectional), the cache would mix tokens across batches
    # and hold large view-retained buffers, causing ~70% VRAM bloat.
    # =============================================================
    use_cache = bool(is_decoder and past_key_value is not None and not is_cross_attention)

    if use_cache:
        self.cache.update(K, V)
        cached_k, cached_v = self.cache.get_cached()
    else:
        # Ensure stale state is dropped so batches stay isolated.
        self.cache.reset()
        cached_k, cached_v = None, None

    if cached_k is not None and cached_v is not None and cached_k.size(0) == (B * H):
        C = cached_k.size(1)
        K_cache = cached_k.view(B, H, C, d)
        V_cache = cached_v.view(B, H, C, d)
        scores_cache = (
            torch.matmul(Q, K_cache.transpose(-2, -1).contiguous())
            * scale
            * torch.sigmoid(torch.clamp(self.sparse_gate, min=-10.0, max=10.0))
        )
        scores = torch.cat(
            [
                scores_lr * torch.sigmoid(torch.clamp(self.LR_scale, min=-10.0, max=10.0)),
                scores_cache,
            ],
            dim=-1,
        )
        V_comb = torch.cat([V, V_cache], dim=2)
    else:
        scores = scores_lr
        V_comb = V

    attn = torch.nan_to_num(F.softmax(scores.float(), dim=-1), nan=0.0).to(dtype)
    attn = self.dropout(attn)

    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)
        elif head_mask.dim() == 3:
            head_mask = head_mask.unsqueeze(-1)
        attn = attn * head_mask

    context = torch.matmul(attn, V_comb)
    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)

    return output, (attn if output_attentions else None)

Must return (output, attention_weights_or_None).

def reset_cache(self) ‑> None

Expand source code

def reset_cache(self) -> None:
    if hasattr(self, "cache") and hasattr(self.cache, "reset"):
        self.cache.reset()

class LoLAAttention (config: Any | None = None,
model_config: Any | None = None,
rank: int | None = None,
lora_rank: int = 16,
dropout: float = 0.1,
comp_config: Any | None = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("lola_attn", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("lola_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("lola", supports_causal=True, supports_gqa=True)
class LoLAAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        rank: int | None = None,
        lora_rank: int = 16,
        dropout: float = 0.1,
        comp_config: Any | None = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config if config is not None else model_config

        def _cfg_get(name: str, fallback: Any) -> Any:
            if cfg is not None:
                val = getattr(cfg, name, None)
                if val is not None:
                    return val
            return fallback

        actual_rank = rank if rank is not None else _cfg_get("rank", 32)
        super().__init__(cfg, actual_rank)

        cfg_head_dim = getattr(cfg, "head_dim", None) if cfg is not None else None
        if cfg is not None:
            self.adapter = ModelAdapter(cfg)
            self.num_heads = self.adapter.num_heads or getattr(cfg, "num_attention_heads", 12)
            self.num_kv_heads = self.adapter.num_key_value_heads or getattr(
                cfg, "num_key_value_heads", self.num_heads
            )
            self.hidden_size = self.adapter.hidden_size or getattr(cfg, "hidden_size", 768)
            self.head_dim = int(
                cfg_head_dim
                or kwargs.get("head_dim")
                or self.adapter.head_dim
                or (self.hidden_size // self.num_heads)
            )
            self.kv_repeat = self.adapter.kv_repeat or 1
        else:
            self.adapter = None
            self.num_heads = kwargs.get("num_heads", kwargs.get("num_attention_heads", 12))
            self.num_kv_heads = kwargs.get("num_key_value_heads", self.num_heads)
            self.hidden_size = kwargs.get("hidden_size", 768)
            self.head_dim = int(kwargs.get("head_dim") or (self.hidden_size // self.num_heads))
            self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))

        _mt = str(getattr(cfg, "model_type", "") or "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        self.is_decoder = (
            (
                getattr(cfg, "is_decoder", False)
                or getattr(cfg, "causal", False)
                or any(_mt.startswith(p) for p in _decoder_prefixes)
                or bool(kwargs.get("is_decoder", False))
            )
            if cfg
            else bool(kwargs.get("is_decoder", False))
        )

        self.all_head_size: int = self.num_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.is_nested_encoder: bool = False

        self.rank: int = actual_rank
        self.lora_rank: int = _cfg_get("lora_rank", lora_rank)

        self.q_proj = LowRankLinear(
            self.hidden_size,
            self.all_head_size,
            self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.k_proj = LowRankLinear(
            self.hidden_size,
            self.kv_head_size,
            self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.v_proj = LowRankLinear(
            self.hidden_size,
            self.kv_head_size,
            self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )
        self.out_proj = LowRankLinear(
            self.all_head_size,
            self.hidden_size,
            self.rank,
            lora_rank=self.lora_rank,
            bias=True,
        )

        self.dropout: nn.Dropout = nn.Dropout(dropout)

        max_seq_len: int = (
            getattr(comp_config, "max_seq_len", 128) if comp_config is not None else 128
        )
        low_rank_cache_size: int = (
            getattr(comp_config, "low_rank_cache_size", 16) if comp_config is not None else 16
        )
        self.cache: LowRankCache = LowRankCache(max_len=max_seq_len, low_rank_k=low_rank_cache_size)
        self.LR_scale: nn.Parameter = nn.Parameter(torch.tensor(4.0))
        self.sparse_gate: nn.Parameter = nn.Parameter(torch.tensor(-3.0))
        self.rotary_emb = (
            getattr(pretrained_attn, "rotary_emb", None) if pretrained_attn is not None else None
        )

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

    def reset_cache(self) -> None:
        if hasattr(self, "cache") and hasattr(self.cache, "reset"):
            self.cache.reset()

    def _set_lowrank_factors(
        self,
        lowrank: LowRankLinear,
        weight: torch.Tensor,
        bias: torch.Tensor | None,
    ) -> None:
        with torch.no_grad():
            W = weight.float()
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            r = min(self.rank, U.size(1), Vt.size(0))
            sqrt_s = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            lowrank.U.data.zero_()
            lowrank.V.data.zero_()
            out_d, in_d = lowrank.U.size(0), lowrank.V.size(1)
            lowrank.U.data[: min(out_d, U.size(0)), :r].copy_(
                (U[: min(out_d, U.size(0)), :r] * sqrt_s).to(lowrank.U.dtype)
            )
            lowrank.V.data[:r, : min(in_d, Vt.size(1))].copy_(
                (sqrt_s.unsqueeze(1) * Vt[:r, : min(in_d, Vt.size(1))]).to(lowrank.V.dtype)
            )
            if lowrank.bias is not None and bias is not None:
                b_val = bias.data if isinstance(bias, nn.Parameter) else bias
                lowrank.bias.data[: min(len(lowrank.bias), len(b_val))].copy_(
                    b_val[: min(len(lowrank.bias), len(b_val))]
                )

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

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        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

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            c_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (weight.dim() == 2 and weight.shape[0] < weight.shape[1])
                ):
                    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:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                p_cls = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__
                b_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or ("Conv1D" in b_cls)
                    or hasattr(c_proj_mod, "nf")
                    or hasattr(base_o, "nf")
                    or (
                        w_o.dim() == 2
                        and (
                            w_o.shape[0] < w_o.shape[1]
                            or (
                                w_o.shape[0] == self.hidden_size
                                and w_o.shape[1] == self.hidden_size
                            )
                        )
                    )
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                extracted["o"] = (w_o, b_o.data.clone() if b_o is not None else None)

        # RoBERTa / BERT nested
        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
            for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]:
                mod = getattr(s, sub_name, None)
                if mod is not None:
                    base_mod = getattr(mod, "base_layer", mod)
                    if hasattr(base_mod, "weight"):
                        extracted[proj_key] = (
                            base_mod.weight.data.clone(),
                            (
                                base_mod.bias.data.clone()
                                if getattr(base_mod, "bias", None) is not None
                                else None
                            ),
                        )

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            dense = src.output.dense
            base_dense = getattr(dense, "base_layer", dense)
            if hasattr(base_dense, "weight"):
                extracted["o"] = (
                    base_dense.weight.data.clone(),
                    (
                        base_dense.bias.data.clone()
                        if getattr(base_dense, "bias", None) is not None
                        else None
                    ),
                )

        # Modern decoders + DistilBERT
        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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)
                        base_l = getattr(layer, "base_layer", layer)
                        if hasattr(base_l, "weight"):
                            extracted[proj] = (
                                base_l.weight.data.clone(),
                                (
                                    base_l.bias.data.clone()
                                    if getattr(base_l, "bias", None) is not None
                                    else None
                                ),
                            )
                            break

        # Nested encoder detection
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            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("Module copy_weights warning. Proceeding with initialized weights.")

        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"LoLAAttention extracted weight for {proj_key} has "
                        f"near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"LoLAAttention extracted weight for output projection has "
                    f"near-zero Frobenius norm ({norm_o.item():.6f})"
                )

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

        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    self._set_lowrank_factors(target, w_oriented, b)
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.U.size(0), target.V.size(1), self.rank)
                    target.U.data.zero_()
                    target.V.data.zero_()
                    target.U.data[:min_d, :min_d].fill_diagonal_(1.0)
                    target.V.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

        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

        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()

        # RoPE
        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"))

        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = d**-0.5

        _cfg_mt = (
            str(getattr(self.config, "model_type", "") or "").lower()
            if self.config is not None
            else ""
        )
        is_decoder = (
            self.is_decoder
            or (getattr(self.config, "is_decoder", False) if self.config is not None else False)
            or any(
                _cfg_mt.startswith(p)
                for p in (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "gemma",
                    "phi",
                    "opt",
                    "bloom",
                    "falcon",
                    "olmo",
                    "mpt",
                    "stablelm",
                )
            )
        )

        scores_lr = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale

        mask = attention_mask if attention_mask is not None else kwargs.get("mask")
        if mask is not None:
            if mask.dim() == 4:
                if mask.dtype == torch.bool:
                    scores_lr = scores_lr.masked_fill(~mask, -10000.0)
                else:
                    scores_lr = scores_lr + mask.to(scores_lr.dtype)
            elif mask.dim() == 3:
                scores_lr = scores_lr + mask.unsqueeze(1).to(scores_lr.dtype)
            elif mask.dim() == 2:
                valid_m = get_token_valid_mask(mask, B, N_k, device)
                if valid_m is not None:
                    scores_lr = scores_lr.masked_fill(
                        (~valid_m).unsqueeze(1).unsqueeze(2), -10000.0
                    )
            else:
                valid_m = get_token_valid_mask(mask, B, N_k, device)
                if valid_m is not None:
                    scores_lr = scores_lr.masked_fill(
                        (~valid_m).unsqueeze(1).unsqueeze(2), -10000.0
                    )

        if is_decoder and N_k >= N_q:
            causal_tril = (
                torch.tril(
                    torch.ones((N_q, N_k), device=device, dtype=torch.bool),
                    diagonal=N_k - N_q,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores_lr = scores_lr.masked_fill(~causal_tril, -10000.0)

        # =============================================================
        # LL-FIX-01/02/03: cache only for streaming decoders with past_key_value.
        # For encoders (bidirectional), the cache would mix tokens across batches
        # and hold large view-retained buffers, causing ~70% VRAM bloat.
        # =============================================================
        use_cache = bool(is_decoder and past_key_value is not None and not is_cross_attention)

        if use_cache:
            self.cache.update(K, V)
            cached_k, cached_v = self.cache.get_cached()
        else:
            # Ensure stale state is dropped so batches stay isolated.
            self.cache.reset()
            cached_k, cached_v = None, None

        if cached_k is not None and cached_v is not None and cached_k.size(0) == (B * H):
            C = cached_k.size(1)
            K_cache = cached_k.view(B, H, C, d)
            V_cache = cached_v.view(B, H, C, d)
            scores_cache = (
                torch.matmul(Q, K_cache.transpose(-2, -1).contiguous())
                * scale
                * torch.sigmoid(torch.clamp(self.sparse_gate, min=-10.0, max=10.0))
            )
            scores = torch.cat(
                [
                    scores_lr * torch.sigmoid(torch.clamp(self.LR_scale, min=-10.0, max=10.0)),
                    scores_cache,
                ],
                dim=-1,
            )
            V_comb = torch.cat([V, V_cache], dim=2)
        else:
            scores = scores_lr
            V_comb = V

        attn = torch.nan_to_num(F.softmax(scores.float(), dim=-1), nan=0.0).to(dtype)
        attn = self.dropout(attn)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn = attn * head_mask

        context = torch.matmul(attn, V_comb)
        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)

        return output, (attn if output_attentions else None)

Every 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.

var supports_rope : bool

The type of the None singleton.

Methods

def reset_cache(self) ‑> None

Expand source code

def reset_cache(self) -> None:
    if hasattr(self, "cache") and hasattr(self.cache, "reset"):
        self.cache.reset()

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class LoRAAttention (config: Any | None = None,
model_config: Any | None = None,
rank: int | None = None,
lora_rank: int = 16,
alpha: float = 32.0,
dropout: float = 0.1,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "lora_attn", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "lora_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class LoRAAttention(BaseAttention):
    """
    LoRA attention wrapper.

    NOTE: This module is deliberately a PEFT swap, NOT a compressor.
    `compression_rate` remains ~1.0; `num_compressed_layers` remains 0.
    Its role is to provide a trainable LoRA residual on top of frozen
    pretrained attention weights.
    """

    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,
        model_config: Any | None = None,
        rank: int | None = None,  # accepted for compat, ignored
        lora_rank: int = 16,
        alpha: float = 32.0,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config if config is not None else model_config
        # LORA2-FIX-01: LoRA rank is what actually matters. `rank` from
        # the sweep (compression rank) is not used for LoRA.
        actual_lora_rank = int(lora_rank) if lora_rank is not None else 16

        super().__init__(cfg, actual_lora_rank)
        # LORA2-FIX-05: cfg may be None
        if cfg is None:
            cfg = getattr(pretrained_attn, "config", None)

        self.adapter = ModelAdapter(cfg)
        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 = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim
        self.kv_repeat = int(self.adapter.kv_repeat)
        self.is_decoder = bool(getattr(cfg, "is_decoder", False)) if cfg is not None else False
        self.is_nested_encoder = False

        self.lora_rank = actual_lora_rank
        self.alpha = float(getattr(cfg, "lora_alpha", alpha)) if cfg is not None else float(alpha)

        self.q_proj = LoRALinear(
            self.hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.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.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.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.alpha,
            freeze_base=True,
        )

        dropout_p = float(
            getattr(
                cfg,
                "dropout",
                getattr(
                    cfg,
                    "hidden_dropout_prob",
                    getattr(cfg, "attention_probs_dropout_prob", dropout),
                ),
            )
            if cfg is not None
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(pretrained_attn, "rotary_emb", None)

        # LORA2-FIX-08: storage accounting
        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 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

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        # Track original bytes
        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. UniversalAdapter
        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 as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # 2. GPT-2 fused c_attn
        if hasattr(src, "c_attn") and any(extracted[p][0] is None for p in ("q", "k", "v")):
            c_attn = src.c_attn
            base_c = getattr(c_attn, "base_layer", c_attn)
            cw = getattr(base_c, "weight", getattr(c_attn, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                is_conv = (
                    c_attn.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (cw.ndim == 2 and cw.shape[0] < cw.shape[1])
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c_attn, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            c_proj = src.c_proj
            base_o = getattr(c_proj, "base_layer", c_proj)
            ow = getattr(base_o, "weight", getattr(c_proj, "weight", None))
            if ow is not None:
                ow = ow.data.clone()
                is_conv_o = (
                    c_proj.__class__.__name__ == "Conv1D"
                    or base_o.__class__.__name__ == "Conv1D"
                    or (ow.ndim == 2 and ow.shape[0] < ow.shape[1])
                )
                if is_conv_o:
                    ow = ow.t().contiguous()
                bo = getattr(base_o, "bias", getattr(c_proj, "bias", None))
                extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # 3. Nested BERT / RoBERTa
        if hasattr(src, "self") and any(extracted[p][0] is None for p in ("q", "k", "v")):
            s = src.self
            if hasattr(s, "query") and extracted["q"][0] is None:
                m = s.query
                extracted["q"] = (
                    m.weight.data,
                    m.bias.data if m.bias is not None else None,
                )
            if hasattr(s, "key") and extracted["k"][0] is None:
                m = s.key
                extracted["k"] = (
                    m.weight.data,
                    m.bias.data if m.bias is not None else None,
                )
            if hasattr(s, "value") and extracted["v"][0] is None:
                m = s.value
                extracted["v"] = (
                    m.weight.data,
                    m.bias.data if m.bias is not None else None,
                )

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            d = src.output.dense
            extracted["o"] = (
                d.weight.data,
                d.bias.data if d.bias is not None else None,
            )

        # 4. Modern decoders and flat modules
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        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"):
                            b = getattr(layer, "bias", None)
                            extracted[proj] = (
                                layer.weight.data,
                                b.data if b is not None else None,
                            )
                            break

        # Detect nested encoder
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            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("LoRAAttention copy_weights: missing %s", missing)

        # Frobenius norm checks
        for k in ("q", "k", "v"):
            w = extracted[k][0]
            if w is not None:
                n = torch.linalg.norm(w.float()).item()
                if n < 1e-4:
                    logger.warning("LoRAAttention: '%s' near-zero norm (%.2e)", k, n)
        if extracted["o"][0] is not None:
            n = torch.linalg.norm(extracted["o"][0].float()).item()
            if n < 1e-4:
                logger.warning("LoRAAttention: 'o' near-zero norm (%.2e)", n)

        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

        with torch.no_grad():
            for key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[key]
                if w is not None:
                    in_dim = self.all_head_size if key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    out_d, in_d = target.weight.shape
                    src_out, src_in = w_oriented.shape
                    # K/V: GQA-safe copy with mean-pooling if needed
                    if key in ("k", "v") and src_out != out_d:
                        if src_out > out_d and src_out % out_d == 0:
                            rep = src_out // out_d
                            w_pooled = w_oriented.view(out_d, rep, -1).mean(dim=1)
                            target.weight.data.copy_(w_pooled)
                            if b is not None and target.bias is not None:
                                bp = b.view(out_d, rep).mean(dim=1)
                                target.bias.data.copy_(bp)
                        elif src_out < out_d and out_d % src_out == 0:
                            rep = out_d // src_out
                            target.weight.data.copy_(w_oriented.repeat(rep, 1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.repeat(rep))
                        else:
                            m0, m1 = min(out_d, src_out), min(in_d, src_in)
                            target.weight.data.zero_()
                            target.weight.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                            if b is not None and target.bias is not None:
                                mb = min(target.bias.shape[0], b.shape[0])
                                target.bias.data.zero_()
                                target.bias.data[:mb].copy_(b[:mb])
                    else:
                        m0, m1 = min(out_d, src_out), min(in_d, src_in)
                        target.weight.data.zero_()
                        target.weight.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                        if b is not None and target.bias is not None:
                            mb = min(target.bias.shape[0], b.shape[0])
                            target.bias.data.zero_()
                            target.bias.data[:mb].copy_(b[:mb])
                elif key == "o" and self.is_nested_encoder:
                    md = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:md, :md].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                target.freeze_base_weights()

        # Track compressed bytes
        try:
            for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

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

    # ------------------------------------------------------------------
    # Mask helper (LORA2-FIX-11)
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_valid_mask(attention_mask, B, N_kv, device):
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                m = attention_mask
            elif attention_mask.dtype.is_floating_point:
                m = attention_mask > -1e4
            else:
                m = attention_mask > 0
            return m if m.shape == (B, N_kv) else None
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            m = v.any(dim=-2).squeeze(1)
            return m if m.shape == (B, N_kv) else None
        return None

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        key_value_states = kwargs.get("key_value_states", encoder_hidden_states)
        is_cross = key_value_states is not None
        kv_input = key_value_states if is_cross else hidden_states

        B, N_q, _ = 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

        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()

        # RoPE
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if (
            position_embeddings is not None
            and not is_cross
            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"))

        # 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
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        # Attention mask
        mask = attention_mask if attention_mask is not None else kwargs.get("mask")
        valid_m = self._normalize_valid_mask(mask, B, N_k, device)
        if valid_m is not None:
            scores = scores.masked_fill(~valid_m[:, None, None, :], -1e4)

        # LORA2-FIX-04: causal only for self-attention
        is_decoder = bool(
            self.is_decoder
            or getattr(self.config, "is_decoder", False)
            or getattr(self.config, "model_type", "")
            in ("gpt2", "qwen", "qwen2", "qwen3_5", "llama")
        )
        if is_decoder and not is_cross:
            causal = torch.tril(torch.ones(N_q, N_k, device=device, dtype=torch.bool))
            scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

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

        if head_mask is not None:
            hm = head_mask.to(attn.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            attn = attn * hm

        context = torch.matmul(attn, V)
        context = context.transpose(1, 2).contiguous().view(B, N_q, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return output, (attn if output_attentions else None)

LoRA attention wrapper.

NOTE: This module is deliberately a PEFT swap, NOT a compressor. compression_rate remains ~1.0; num_compressed_layers remains 0. Its role is to provide a trainable LoRA residual on top of frozen pretrained attention weights.

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

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

Expand source code

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

    Base path:
        y_base = x @ W^T + b  (frozen when freeze_base=True)
    LoRA path:
        y_lora = (x @ lora_A^T) @ lora_B^T * scaling  (trainable)
    Total output:
        y = y_base + y_lora

    Uses standard `lora_A` and `lora_B` parameter naming recognized by framework recovery optimizers.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        freeze_base: bool = False,
        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:
        """Freeze base projection parameters so only LoRA parameters adapt."""
        self.weight.requires_grad = False
        if self.bias is not None:
            self.bias.requires_grad = False
        if self.lora_rank > 0 and self.lora_A is not None:
            self.lora_A.requires_grad = True
        if self.lora_rank > 0 and self.lora_B is not None:
            self.lora_B.requires_grad = True

    def unfreeze_base_weights(self) -> None:
        """Unfreeze base projection parameters."""
        self.weight.requires_grad = True
        if self.bias is not None:
            self.bias.requires_grad = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.weight.to(x.dtype) if self.weight.dtype != x.dtype else self.weight
        b = (
            self.bias.to(x.dtype)
            if (self.bias is not None and self.bias.dtype != x.dtype)
            else self.bias
        )
        out = F.linear(x, w, b)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
            bb = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
            lora_out = F.linear(F.linear(x, a), bb) * self.lora_scaling
            out = out + lora_out
        return out

Linear projection layer equipped with frozen base weights and trainable residual LoRA adapters.

Base path: y_base = x @ W^T + b (frozen when freeze_base=True) LoRA path: y_lora = (x @ lora_A^T) @ lora_B^T * scaling (trainable) Total output: y = y_base + y_lora

Uses standard lora_A and lora_B parameter naming recognized by framework recovery optimizers.

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:
    w = self.weight.to(x.dtype) if self.weight.dtype != x.dtype else self.weight
    b = (
        self.bias.to(x.dtype)
        if (self.bias is not None and self.bias.dtype != x.dtype)
        else self.bias
    )
    out = F.linear(x, w, b)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A
        bb = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B
        lora_out = F.linear(F.linear(x, a), bb) * self.lora_scaling
        out = out + lora_out
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    """Freeze base projection parameters so only LoRA parameters adapt."""
    self.weight.requires_grad = False
    if self.bias is not None:
        self.bias.requires_grad = False
    if self.lora_rank > 0 and self.lora_A is not None:
        self.lora_A.requires_grad = True
    if self.lora_rank > 0 and self.lora_B is not None:
        self.lora_B.requires_grad = True

Freeze base projection parameters so only LoRA parameters adapt.

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

def unfreeze_base_weights(self) ‑> None

Expand source code

def unfreeze_base_weights(self) -> None:
    """Unfreeze base projection parameters."""
    self.weight.requires_grad = True
    if self.bias is not None:
        self.bias.requires_grad = True

Unfreeze base projection parameters.

class LowRankCache (max_len: int = 128, low_rank_k: int = 16)

Expand source code

class LowRankCache(nn.Module):
    def __init__(self, max_len: int = 128, low_rank_k: int = 16) -> None:
        super().__init__()
        self.max_len: int = max_len
        self.low_rank_k: int = low_rank_k
        self.register_buffer("recent_k", torch.zeros(0), persistent=False)
        self.recent_k: torch.Tensor
        self.register_buffer("recent_v", torch.zeros(0), persistent=False)
        self.recent_v: torch.Tensor
        self.register_buffer("sparse_k", torch.zeros(0), persistent=False)
        self.sparse_k: torch.Tensor
        self.register_buffer("sparse_v", torch.zeros(0), persistent=False)
        self.sparse_v: torch.Tensor
        self.register_buffer("low_rank_k_buf", torch.zeros(0), persistent=False)
        self.low_rank_k_buf: torch.Tensor
        self.register_buffer("low_rank_v_buf", torch.zeros(0), persistent=False)
        self.low_rank_v_buf: torch.Tensor
        self.register_buffer("_cache_proj", torch.zeros(0), persistent=False)
        self._cache_proj: torch.Tensor
        self._update_step: int = 0

    def reset(self) -> None:
        self.recent_k = torch.zeros(0, device=self.recent_k.device, dtype=self.recent_k.dtype)
        self.recent_v = torch.zeros(0, device=self.recent_v.device, dtype=self.recent_v.dtype)
        self.sparse_k = torch.zeros(0, device=self.sparse_k.device, dtype=self.sparse_k.dtype)
        self.sparse_v = torch.zeros(0, device=self.sparse_v.device, dtype=self.sparse_v.dtype)
        self.low_rank_k_buf = torch.zeros(
            0, device=self.low_rank_k_buf.device, dtype=self.low_rank_k_buf.dtype
        )
        self.low_rank_v_buf = torch.zeros(
            0, device=self.low_rank_v_buf.device, dtype=self.low_rank_v_buf.dtype
        )
        self._update_step = 0

    def update(self, K: torch.Tensor, V: torch.Tensor) -> None:
        B, H, N, D = K.shape
        k_flat = K.detach().reshape(B * H, N, D)
        v_flat = V.detach().reshape(B * H, N, D)

        # Recent cache
        if (
            self.recent_k.numel() == 0
            or self.recent_k.size(0) != B * H
            or self.recent_k.size(-1) != D
        ):
            self.recent_k = k_flat[:, -self.max_len :, :].contiguous().clone()
            self.recent_v = v_flat[:, -self.max_len :, :].contiguous().clone()
        else:
            self.recent_k = (
                torch.cat([self.recent_k, k_flat], dim=1)[:, -self.max_len :, :]
                .contiguous()
                .clone()
            )
            self.recent_v = (
                torch.cat([self.recent_v, v_flat], dim=1)[:, -self.max_len :, :]
                .contiguous()
                .clone()
            )

        # Sparse cache (sample every 4th token)
        sparse_idx = torch.arange(0, N, 4, device=K.device)
        sparse_k = k_flat[:, sparse_idx, :]
        sparse_v = v_flat[:, sparse_idx, :]
        if (
            self.sparse_k.numel() == 0
            or self.sparse_k.size(0) != B * H
            or self.sparse_k.size(-1) != D
        ):
            self.sparse_k = sparse_k[:, -self.max_len // 2 :, :].contiguous().clone()
            self.sparse_v = sparse_v[:, -self.max_len // 2 :, :].contiguous().clone()
        else:
            self.sparse_k = (
                torch.cat([self.sparse_k, sparse_k], dim=1)[:, -self.max_len // 2 :, :]
                .contiguous()
                .clone()
            )
            self.sparse_v = (
                torch.cat([self.sparse_v, sparse_v], dim=1)[:, -self.max_len // 2 :, :]
                .contiguous()
                .clone()
            )

        # Low-rank subspace cache
        if self.recent_k.size(1) >= self.low_rank_k:
            step = getattr(self, "_update_step", 0) + 1
            self._update_step = step
            if (
                step % 50 == 1
                or self.low_rank_k_buf.numel() == 0
                or self.low_rank_k_buf.size(0) != B * H
            ):
                with torch.no_grad():
                    rk = self.recent_k.float()
                    rv = self.recent_v.float()
                    if (
                        self._cache_proj.numel() == 0
                        or self._cache_proj.size(0) != D
                        or self._cache_proj.size(1) != self.low_rank_k
                        or self._cache_proj.device != rk.device
                    ):
                        proj = torch.empty(
                            D, self.low_rank_k, device=rk.device, dtype=torch.float32
                        )
                        nn.init.orthogonal_(proj)
                        self._cache_proj = proj

                    try:
                        y_k = torch.matmul(rk, self._cache_proj)
                        q_orth_k, _ = torch.linalg.qr(y_k)
                        self.low_rank_k_buf = torch.matmul(
                            q_orth_k,
                            torch.matmul(q_orth_k.transpose(-2, -1).contiguous(), rk),
                        ).to(self.recent_k.dtype)

                        y_v = torch.matmul(rv, self._cache_proj)
                        q_orth_v, _ = torch.linalg.qr(y_v)
                        self.low_rank_v_buf = torch.matmul(
                            q_orth_v,
                            torch.matmul(q_orth_v.transpose(-2, -1).contiguous(), rv),
                        ).to(self.recent_v.dtype)
                    except Exception:
                        pass

    def get_cached(self) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        if self.recent_k.numel() == 0:
            return None, None
        B_H = self.recent_k.size(0)
        tensors_k = [self.recent_k]
        tensors_v = [self.recent_v]
        if self.sparse_k.numel() > 0 and self.sparse_k.size(0) == B_H:
            tensors_k.append(self.sparse_k)
            tensors_v.append(self.sparse_v)
        if self.low_rank_k_buf.numel() > 0 and self.low_rank_k_buf.size(0) == B_H:
            tensors_k.append(self.low_rank_k_buf)
            tensors_v.append(self.low_rank_v_buf)
        return torch.cat(tensors_k, dim=1), torch.cat(tensors_v, dim=1)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Methods

def get_cached(self) ‑> tuple[torch.Tensor | None, torch.Tensor | None]

Expand source code

def get_cached(self) -> tuple[torch.Tensor | None, torch.Tensor | None]:
    if self.recent_k.numel() == 0:
        return None, None
    B_H = self.recent_k.size(0)
    tensors_k = [self.recent_k]
    tensors_v = [self.recent_v]
    if self.sparse_k.numel() > 0 and self.sparse_k.size(0) == B_H:
        tensors_k.append(self.sparse_k)
        tensors_v.append(self.sparse_v)
    if self.low_rank_k_buf.numel() > 0 and self.low_rank_k_buf.size(0) == B_H:
        tensors_k.append(self.low_rank_k_buf)
        tensors_v.append(self.low_rank_v_buf)
    return torch.cat(tensors_k, dim=1), torch.cat(tensors_v, dim=1)

def reset(self) ‑> None

Expand source code

def reset(self) -> None:
    self.recent_k = torch.zeros(0, device=self.recent_k.device, dtype=self.recent_k.dtype)
    self.recent_v = torch.zeros(0, device=self.recent_v.device, dtype=self.recent_v.dtype)
    self.sparse_k = torch.zeros(0, device=self.sparse_k.device, dtype=self.sparse_k.dtype)
    self.sparse_v = torch.zeros(0, device=self.sparse_v.device, dtype=self.sparse_v.dtype)
    self.low_rank_k_buf = torch.zeros(
        0, device=self.low_rank_k_buf.device, dtype=self.low_rank_k_buf.dtype
    )
    self.low_rank_v_buf = torch.zeros(
        0, device=self.low_rank_v_buf.device, dtype=self.low_rank_v_buf.dtype
    )
    self._update_step = 0

def update(self, K: torch.Tensor, V: torch.Tensor) ‑> None

Expand source code

def update(self, K: torch.Tensor, V: torch.Tensor) -> None:
    B, H, N, D = K.shape
    k_flat = K.detach().reshape(B * H, N, D)
    v_flat = V.detach().reshape(B * H, N, D)

    # Recent cache
    if (
        self.recent_k.numel() == 0
        or self.recent_k.size(0) != B * H
        or self.recent_k.size(-1) != D
    ):
        self.recent_k = k_flat[:, -self.max_len :, :].contiguous().clone()
        self.recent_v = v_flat[:, -self.max_len :, :].contiguous().clone()
    else:
        self.recent_k = (
            torch.cat([self.recent_k, k_flat], dim=1)[:, -self.max_len :, :]
            .contiguous()
            .clone()
        )
        self.recent_v = (
            torch.cat([self.recent_v, v_flat], dim=1)[:, -self.max_len :, :]
            .contiguous()
            .clone()
        )

    # Sparse cache (sample every 4th token)
    sparse_idx = torch.arange(0, N, 4, device=K.device)
    sparse_k = k_flat[:, sparse_idx, :]
    sparse_v = v_flat[:, sparse_idx, :]
    if (
        self.sparse_k.numel() == 0
        or self.sparse_k.size(0) != B * H
        or self.sparse_k.size(-1) != D
    ):
        self.sparse_k = sparse_k[:, -self.max_len // 2 :, :].contiguous().clone()
        self.sparse_v = sparse_v[:, -self.max_len // 2 :, :].contiguous().clone()
    else:
        self.sparse_k = (
            torch.cat([self.sparse_k, sparse_k], dim=1)[:, -self.max_len // 2 :, :]
            .contiguous()
            .clone()
        )
        self.sparse_v = (
            torch.cat([self.sparse_v, sparse_v], dim=1)[:, -self.max_len // 2 :, :]
            .contiguous()
            .clone()
        )

    # Low-rank subspace cache
    if self.recent_k.size(1) >= self.low_rank_k:
        step = getattr(self, "_update_step", 0) + 1
        self._update_step = step
        if (
            step % 50 == 1
            or self.low_rank_k_buf.numel() == 0
            or self.low_rank_k_buf.size(0) != B * H
        ):
            with torch.no_grad():
                rk = self.recent_k.float()
                rv = self.recent_v.float()
                if (
                    self._cache_proj.numel() == 0
                    or self._cache_proj.size(0) != D
                    or self._cache_proj.size(1) != self.low_rank_k
                    or self._cache_proj.device != rk.device
                ):
                    proj = torch.empty(
                        D, self.low_rank_k, device=rk.device, dtype=torch.float32
                    )
                    nn.init.orthogonal_(proj)
                    self._cache_proj = proj

                try:
                    y_k = torch.matmul(rk, self._cache_proj)
                    q_orth_k, _ = torch.linalg.qr(y_k)
                    self.low_rank_k_buf = torch.matmul(
                        q_orth_k,
                        torch.matmul(q_orth_k.transpose(-2, -1).contiguous(), rk),
                    ).to(self.recent_k.dtype)

                    y_v = torch.matmul(rv, self._cache_proj)
                    q_orth_v, _ = torch.linalg.qr(y_v)
                    self.low_rank_v_buf = torch.matmul(
                        q_orth_v,
                        torch.matmul(q_orth_v.transpose(-2, -1).contiguous(), rv),
                    ).to(self.recent_v.dtype)
                except Exception:
                    pass

class LowRankLinear (in_features: int,
out_features: int,
rank: int,
lora_rank: int = 16,
bias: bool = True)

Expand source code

class LowRankLinear(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        lora_rank: int = 16,
        bias: bool = True,
    ) -> None:
        super().__init__()
        self.in_features: int = in_features
        self.out_features: int = out_features
        self.rank: int = rank
        self.lora_rank: int = min(lora_rank, min(in_features, out_features))
        self.lora_scaling: float = (32.0 / self.lora_rank) if self.lora_rank > 0 else 1.0

        self.U: nn.Parameter = nn.Parameter(torch.empty(out_features, rank), requires_grad=False)
        self.V: nn.Parameter = nn.Parameter(torch.empty(rank, in_features), requires_grad=False)
        if bias:
            self.bias: nn.Parameter | None = nn.Parameter(
                torch.empty(out_features), requires_grad=False
            )
        else:
            self.register_parameter("bias", None)

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

        self.reset_parameters()

    def reset_parameters(self) -> None:
        with torch.no_grad():
            u_fp32 = torch.empty_like(self.U, dtype=torch.float32)
            v_fp32 = torch.empty_like(self.V, dtype=torch.float32)
            nn.init.orthogonal_(u_fp32)
            nn.init.orthogonal_(v_fp32)
            self.U.data.copy_(u_fp32.to(self.U.dtype))
            self.V.data.copy_(v_fp32.to(self.V.dtype))
            if self.bias is not None:
                nn.init.zeros_(self.bias)

        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

    @property
    def weight(self) -> torch.Tensor:
        w = torch.matmul(self.U, self.V)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            w = w + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
        return w

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

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    w = torch.matmul(self.U, self.V)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        w = w + torch.matmul(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:
    Vx = F.linear(x, self.V)
    out = F.linear(Vx, self.U)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
        out = out + lora_out
    if self.bias is not None:
        out = out + self.bias
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    with torch.no_grad():
        u_fp32 = torch.empty_like(self.U, dtype=torch.float32)
        v_fp32 = torch.empty_like(self.V, dtype=torch.float32)
        nn.init.orthogonal_(u_fp32)
        nn.init.orthogonal_(v_fp32)
        self.U.data.copy_(u_fp32.to(self.U.dtype))
        self.V.data.copy_(v_fp32.to(self.V.dtype))
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    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 LowRankSparseLinear (down: nn.Linear,
up: nn.Linear,
sparse_weight: torch.Tensor,
bias: torch.Tensor | None = None)

Expand source code

class LowRankSparseLinear(nn.Module):
    """
    y = up(down(x)) + F.linear(x, S) + b

    S is a trainable sparse residual. Its L1 gradient encourages zero
    entries; periodic `prune_sparse` enforces hard threshold.

    SP-FIX-02: bias stored as a proper nn.Parameter.
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        down: nn.Linear,
        up: nn.Linear,
        sparse_weight: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> None:
        super().__init__()
        self.down = down
        self.up = up
        self.sparse_weight = nn.Parameter(sparse_weight)
        if bias is not None:
            self.bias = nn.Parameter(bias.clone())
        else:
            self.register_parameter("bias", None)

        self.in_features = down.in_features
        self.out_features = up.out_features

    @property
    def weight(self) -> torch.Tensor:
        return torch.matmul(self.up.weight, self.down.weight) + self.sparse_weight

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SP-FIX-11: dtype-safe
        down_w = self.down.weight.to(x.dtype)
        up_w = self.up.weight.to(x.dtype)
        up_b = self.up.bias.to(x.dtype) if self.up.bias is not None else None

        x_lr = F.linear(F.linear(x, down_w), up_w, up_b)
        x_sp = F.linear(x, self.sparse_weight.to(x.dtype))
        out = x_lr + x_sp
        if self.bias is not None:
            out = out + self.bias.to(out.dtype)
        return out

    def prune_sparse(self, sparsity: float) -> None:
        """Hard threshold: keep only top (1 - sparsity) fraction of entries."""
        with torch.no_grad():
            w = self.sparse_weight.data
            if w.numel() == 0:
                return
            keep = int(w.numel() * (1.0 - sparsity))
            if keep <= 0:
                w.zero_()
                return
            if keep >= w.numel():
                return
            flat = w.abs().view(-1)
            kth = torch.kthvalue(flat, w.numel() - keep).values
            w.mul_(w.abs() >= kth)

    def sparse_l1_loss(self) -> torch.Tensor:
        return self.sparse_weight.abs().sum()

y = up(down(x)) + F.linear(x, S) + b

S is a trainable sparse residual. Its L1 gradient encourages zero entries; periodic prune_sparse enforces hard threshold.

SP-FIX-02: bias stored as a proper nn.Parameter.

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

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    return torch.matmul(self.up.weight, self.down.weight) + self.sparse_weight

Methods

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

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # SP-FIX-11: dtype-safe
    down_w = self.down.weight.to(x.dtype)
    up_w = self.up.weight.to(x.dtype)
    up_b = self.up.bias.to(x.dtype) if self.up.bias is not None else None

    x_lr = F.linear(F.linear(x, down_w), up_w, up_b)
    x_sp = F.linear(x, self.sparse_weight.to(x.dtype))
    out = x_lr + x_sp
    if self.bias is not None:
        out = out + self.bias.to(out.dtype)
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

def prune_sparse(self, sparsity: float) ‑> None

Expand source code

def prune_sparse(self, sparsity: float) -> None:
    """Hard threshold: keep only top (1 - sparsity) fraction of entries."""
    with torch.no_grad():
        w = self.sparse_weight.data
        if w.numel() == 0:
            return
        keep = int(w.numel() * (1.0 - sparsity))
        if keep <= 0:
            w.zero_()
            return
        if keep >= w.numel():
            return
        flat = w.abs().view(-1)
        kth = torch.kthvalue(flat, w.numel() - keep).values
        w.mul_(w.abs() >= kth)

Hard threshold: keep only top (1 - sparsity) fraction of entries.

def sparse_l1_loss(self) ‑> torch.Tensor

Expand source code

def sparse_l1_loss(self) -> torch.Tensor:
    return self.sparse_weight.abs().sum()

class MLAAttention (config: Any | None = None,
rank: int | None = 128,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
max_error: float = 0.35,
fallback_to_original: bool = True,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("mla", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("mla_attention", supports_causal=True, supports_gqa=True)
class MLAAttention(BaseAttention):
    """
    Multi-Head Latent Attention (MLA) — DeepSeek-V2 style.

    Idea: compress K and V into a low-rank latent `c_kv` of size `rank`,
    then up-project back to per-head K and V.

        c_kv = kv_down(x)             # (B, N, rank)
        k    = k_up(c_kv)              # (B, N, num_kv_heads * head_dim)
        v    = v_up(c_kv)              # (B, N, num_kv_heads * head_dim)

    Query is projected with a LoRA-augmented linear. K/V down/up projections
    are initialized via joint SVD of [W_k; W_v] and frozen.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "low_rank"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 128,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        max_error: float = 0.35,
        fallback_to_original: bool = True,
        **kwargs: Any,
    ) -> None:
        # MLA-FIX-05: config None-safe
        if config is None:

            class _Cfg:
                pass

            config = _Cfg()
            config.hidden_size = kwargs.get("hidden_size", 768)
            config.num_attention_heads = kwargs.get("num_attention_heads", 12)
            config.head_dim = kwargs.get("head_dim", 64)
            config.model_type = kwargs.get("model_type", "")
            config.num_key_value_heads = kwargs.get(
                "num_key_value_heads", kwargs.get("num_attention_heads", 12)
            )
            config.is_decoder = kwargs.get("is_decoder", False)

        try:
            super().__init__(config, rank if rank is not None else 128)
        except TypeError:
            super().__init__(config)

        self.hidden_size: int = int(
            getattr(
                config,
                "hidden_size",
                getattr(config, "d_model", kwargs.get("hidden_size", 768)),
            )
        )
        self.num_heads: int = int(
            getattr(
                config,
                "num_attention_heads",
                getattr(
                    config,
                    "n_heads",
                    getattr(config, "n_head", getattr(config, "num_heads", 12)),
                ),
            )
        )
        self.head_dim: int = int(
            getattr(config, "head_dim", self.hidden_size // max(1, self.num_heads))
        )
        self.q_dim: int = int(getattr(config, "q_proj_out", self.num_heads * self.head_dim))
        self.all_head_size: int = self.q_dim

        cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_dim: int = self.num_kv_heads * self.head_dim
        self.kv_head_size: int = self.kv_dim
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))

        self.rank: int = int(getattr(config, "mla_rank", rank if rank is not None else 128))
        self.lora_rank: int = int(lora_rank)
        self.lora_alpha: float = float(lora_alpha)
        self.max_error: float = float(max_error)
        self.fallback_to_original: bool = bool(fallback_to_original)

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(getattr(config, "model_type", "") or "").lower()
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            self.is_decoder = (
                bool(getattr(config, "is_decoder", False))
                or bool(getattr(config, "causal", False))
                or any(_mt.startswith(p) for p in _dec_pfx)
            )

        self.is_nested_encoder: bool = False

        # ---- projections ----
        self.q_proj = LoRALinear(
            self.hidden_size,
            self.q_dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.kv_down = nn.Linear(self.hidden_size, self.rank, bias=False)
        self.k_up = nn.Linear(self.rank, self.kv_dim, bias=False)
        self.v_up = nn.Linear(self.rank, self.kv_dim, bias=False)
        self.out_proj = LoRALinear(
            self.q_dim,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
        )
        self.dropout = nn.Dropout(dropout_p)

        # ---- storage accounting (MLA-FIX-08) ----
        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 0
        self._reconstruction_error_kv: float | None = None

        # ---- source attention (for injector) ----
        self._source_attn_id: int | None = None
        self.triton_kwargs = {}

        pretrained_attn = kwargs.get("pretrained_attn")
        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    # ------------------------------------------------------------------
    # Aliases (MLA-FIX-02, MLA-FIX-10)
    # ------------------------------------------------------------------
    @property
    def query(self):
        return self.q_proj

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

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

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

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

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

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

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

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

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

    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        # MLA-FIX-06: idempotency guard
        if self._source_attn_id is not None and id(src) == self._source_attn_id:
            logger.debug("MLAAttention.copy_weights: source already applied, skipping")
            return

        # --- extract weights ---
        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        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 as e:
                logger.debug("UniversalAdapter fallback failed for %s: %s", proj, e)

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            c_attn = src.c_attn
            base_c = getattr(c_attn, "base_layer", c_attn)
            weight = getattr(base_c, "weight", getattr(c_attn, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                is_conv = (
                    c_attn.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
                )
                if is_conv:
                    weight = weight.t().contiguous()
                split = weight.shape[0] // 3
                bias = getattr(base_c, "bias", getattr(c_attn, "bias", None))
                bd = bias.data if bias is not None else None
                extracted["q"] = (
                    weight[:split],
                    bd[:split] if bd is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split],
                    bd[split : 2 * split] if bd is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :],
                    bd[2 * split :] if bd is not None else None,
                )

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            c_proj = src.c_proj
            base_o = getattr(c_proj, "base_layer", c_proj)
            w_o = getattr(base_o, "weight", getattr(c_proj, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                if (
                    w_o.dim() == 2
                    and w_o.shape[0] != self.hidden_size
                    and w_o.shape[1] == self.hidden_size
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj, "bias", None))
                extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # BERT / RoBERTa nested
        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") and extracted["q"][0] is None:
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key") and extracted["k"][0] is None:
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value") and extracted["v"][0] is None:
                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))

        # Flat modules
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        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

        # --- nested encoder detection (MLA-FIX-07) ---
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or cls_name in ("BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention")
            or cls_name.endswith("SelfAttention")
            and not has_out_proj
            or cls_name.endswith("Attention")
            and not has_out_proj
        ):
            is_nested = True
        self.is_nested_encoder = bool(is_nested)

        # --- required weights check ---
        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("MLAAttention.copy_weights: missing %s; aborting.", missing)
            return

        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

        # --- storage accounting ---
        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        with torch.no_grad():
            # ---------- 1. Q ----------
            w_q, b_q = extracted["q"]
            w_q_oriented = _orient(w_q, self.hidden_size)
            out_q = min(self.q_dim, w_q_oriented.shape[0])
            in_q = min(self.hidden_size, w_q_oriented.shape[1])
            self.q_proj.weight.data.zero_()
            self.q_proj.weight.data[:out_q, :in_q].copy_(
                w_q_oriented[:out_q, :in_q].to(self.q_proj.weight.device, self.q_proj.weight.dtype)
            )
            if b_q is not None and self.q_proj.bias is not None:
                mb = min(self.q_proj.bias.numel(), b_q.numel())
                self.q_proj.bias.data.zero_()
                self.q_proj.bias.data[:mb].copy_(
                    b_q[:mb].to(self.q_proj.bias.device, self.q_proj.bias.dtype)
                )
            self.q_proj.freeze_base_weights()

            # ---------- 2. Output ----------
            w_o, b_o = extracted["o"]
            if w_o is not None:
                w_o_oriented = _orient(w_o, self.q_dim)
                out_o = min(self.hidden_size, w_o_oriented.shape[0])
                in_o = min(self.q_dim, w_o_oriented.shape[1])
                self.out_proj.weight.data.zero_()
                self.out_proj.weight.data[:out_o, :in_o].copy_(
                    w_o_oriented[:out_o, :in_o].to(
                        self.out_proj.weight.device, self.out_proj.weight.dtype
                    )
                )
                if b_o is not None and self.out_proj.bias is not None:
                    mb = min(self.out_proj.bias.numel(), b_o.numel())
                    self.out_proj.bias.data.zero_()
                    self.out_proj.bias.data[:mb].copy_(
                        b_o[:mb].to(self.out_proj.bias.device, self.out_proj.bias.dtype)
                    )
            elif self.is_nested_encoder:
                # Identity init for nested encoders
                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_()
            self.out_proj.freeze_base_weights()

            # ---------- 3. K/V low-rank SVD ----------
            w_k, _ = extracted["k"]
            w_v, _ = extracted["v"]
            w_k_f = _orient(w_k, self.hidden_size).float()
            w_v_f = _orient(w_v, self.hidden_size).float()

            # GQA dimension matching
            def _match_gqa(w, target_dim, hidden_size):
                if w.shape[0] < target_dim and target_dim % w.shape[0] == 0:
                    rep = target_dim // w.shape[0]
                    return w.repeat(rep, 1)
                if w.shape[0] > target_dim and w.shape[0] % target_dim == 0:
                    rep = w.shape[0] // target_dim
                    return w.view(target_dim, rep, -1).mean(dim=1)
                return w

            w_k_f = _match_gqa(w_k_f, self.kv_dim, self.hidden_size)
            w_v_f = _match_gqa(w_v_f, self.kv_dim, self.hidden_size)

            # Joint SVD of [K; V]
            w_kv = torch.cat(
                [
                    w_k_f[: self.kv_dim, : self.hidden_size],
                    w_v_f[: self.kv_dim, : self.hidden_size],
                ],
                dim=0,
            )  # (2 * kv_dim, hidden_size)

            u_kv, s_kv, vh_kv = torch.linalg.svd(w_kv, full_matrices=False)

            r = min(
                self.rank,
                s_kv.shape[0],
                self.kv_down.weight.shape[0],
                self.k_up.weight.shape[1],
                self.v_up.weight.shape[1],
            )

            # MLA-FIX-01: NO norm_scale. The SVD low-rank approximation
            # naturally has the correct norm: ||U_r Σ_r V_r^T||_F = ||Σ_r||_F.
            self.kv_down.weight.data.zero_()
            self.k_up.weight.data.zero_()
            self.v_up.weight.data.zero_()

            # kv_down: (rank, hidden_size) = vh_kv[:r, :]
            self.kv_down.weight.data[:r, :].copy_(
                vh_kv[:r, :].to(self.kv_down.weight.device, self.kv_down.weight.dtype)
            )

            # up: (2 * kv_dim, r) = U[:, :r] * Σ[:r]
            up = (u_kv[:, :r] * s_kv[:r].unsqueeze(0)).to(
                self.k_up.weight.device, self.k_up.weight.dtype
            )

            d_k = min(self.kv_dim, self.k_up.weight.shape[0])
            d_v = min(self.kv_dim, self.v_up.weight.shape[0])
            # K rows come first, V rows next
            self.k_up.weight.data[:d_k, :r].copy_(up[:d_k, :r])
            self.v_up.weight.data[:d_v, :r].copy_(up[self.kv_dim : self.kv_dim + d_v, :r])

            self.kv_down.weight.requires_grad = False
            self.k_up.weight.requires_grad = False
            self.v_up.weight.requires_grad = False

            # MLA-FIX-04: reconstruction error validation
            W_recon = torch.cat(
                [
                    self.k_up.weight.data.float() @ self.kv_down.weight.data.float(),
                    self.v_up.weight.data.float() @ self.kv_down.weight.data.float(),
                ],
                dim=0,
            )
            num = torch.norm(w_kv - W_recon, p="fro").item()
            den = torch.norm(w_kv, p="fro").item() + 1e-12
            self._reconstruction_error_kv = num / den

            if self._reconstruction_error_kv > self.max_error:
                logger.warning(
                    "MLAAttention: K/V reconstruction error %.4f > max_error %.4f. "
                    "Module marked for fallback.",
                    self._reconstruction_error_kv,
                    self.max_error,
                )
            else:
                logger.info(
                    "MLAAttention: K/V reconstruction error %.4f (max allowed %.4f).",
                    self._reconstruction_error_kv,
                    self.max_error,
                )

        # ---- storage accounting (compressed) ----
        try:
            for m in (self.q_proj, self.kv_down, self.k_up, self.v_up, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # ---- is_decoder / nested override ----
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)
        elif (
            cls_name in ("GPT2Attention", "GPT2Block")
            or hasattr(src, "c_attn")
            or "gpt" in cls_name.lower()
        ):
            self.is_decoder = True

        # Record source id for idempotency
        self._source_attn_id = id(src)

    # ------------------------------------------------------------------
    def is_compatible(self) -> bool:
        """Return True if reconstruction error is within tolerance."""
        if self._reconstruction_error_kv is None:
            return True
        return self._reconstruction_error_kv <= self.max_error

    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            return p.numel() * p.element_size() if isinstance(p, torch.Tensor) else 0

        return {
            "orig": self._orig_storage_bytes,
            "comp": self._comp_storage_bytes,
            "kv_reconstruction_error": self._reconstruction_error_kv,
        }

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ):
        B, N, _ = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

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

        # ---- projections ----
        q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k_raw = triton_dispatch(self.triton_kernel, kv_input, self.k_up.weight, self.kv_down.weight)
        v_raw = triton_dispatch(self.triton_kernel, kv_input, self.v_up.weight, self.kv_down.weight)
        k = k_raw.view(B, kv_len, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        v = v_raw.view(B, kv_len, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()

        # ---- RoPE ----
        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
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # ---- GQA expansion ----
        if self.num_kv_heads < self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        scale = 1.0 / math.sqrt(self.head_dim)
        eff_mask = encoder_attention_mask if is_cross else attention_mask

        # ---- SDPA fast path ----
        if eff_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 = bool(self.is_decoder and not is_cross and N > 1)
                out = F.scaled_dot_product_attention(
                    _qs,
                    _ks,
                    _vs,
                    is_causal=_is_causal,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                out = out.to(dtype).transpose(1, 2).contiguous().view(B, N, self.all_head_size)
                out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
                if self.is_nested_encoder:
                    return out, None
                return self.out_proj(out), None
            except Exception:
                pass

        # ---- Manual path ----
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross and N > 1:
            diag = 1 + max(0, kv_len - N)
            causal = (
                torch.triu(
                    torch.full((N, kv_len), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=diag,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        if eff_mask is not None:
            if eff_mask.dim() == 2:
                scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                scores = scores + eff_mask.float()

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

        if head_mask is not None:
            hm = head_mask.to(attn.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            attn = attn * hm

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

        out = context if self.is_nested_encoder else self.out_proj(context)
        return out, (attn if output_attentions else None)

Multi-Head Latent Attention (MLA) — DeepSeek-V2 style.

Idea: compress K and V into a low-rank latent c_kv of size rank, then up-project back to per-head K and V.

c_kv = kv_down(x)             # (B, N, rank)
k    = k_up(c_kv)              # (B, N, num_kv_heads * head_dim)
v    = v_up(c_kv)              # (B, N, num_kv_heads * head_dim)

Query is projected with a LoRA-augmented linear. K/V down/up projections are initialized via joint SVD of [W_k; W_v] and frozen.

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.

var triton_kernel : str

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_up

prop key

Expand source code

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

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_up

prop value

Expand source code

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

Methods

def is_compatible(self) ‑> bool

Expand source code

def is_compatible(self) -> bool:
    """Return True if reconstruction error is within tolerance."""
    if self._reconstruction_error_kv is None:
        return True
    return self._reconstruction_error_kv <= self.max_error

Return True if reconstruction error is within tolerance.

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        return p.numel() * p.element_size() if isinstance(p, torch.Tensor) else 0

    return {
        "orig": self._orig_storage_bytes,
        "comp": self._comp_storage_bytes,
        "kv_reconstruction_error": self._reconstruction_error_kv,
    }

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class Mamba (d_model: int = 768,
**kwargs: Any)

Expand source code

class Mamba(nn.Module):
    def __init__(self, d_model: int = 768, **kwargs: Any) -> None:
        super().__init__()
        self.in_proj = nn.Linear(d_model, 2 * d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        nn.init.zeros_(self.out_proj.weight)
        if self.out_proj.bias is not None:
            nn.init.zeros_(self.out_proj.bias)

    def forward(self, x: torch.Tensor, **kwargs: Any) -> torch.Tensor:
        x1, x2 = self.in_proj(x).chunk(2, dim=-1)
        y = F.silu(x1) * x2
        return self.out_proj(y)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

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

Ancestors

• torch.nn.modules.module.Module

Methods

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

Expand source code

def forward(self, x: torch.Tensor, **kwargs: Any) -> torch.Tensor:
    x1, x2 = self.in_proj(x).chunk(2, dim=-1)
    y = F.silu(x1) * x2
    return self.out_proj(y)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class Mamba2 (d_model: int,
d_state: int = 16,
d_conv: int = 4,
expand: int = 2,
**kwargs: Any)

Expand source code

class Mamba2(nn.Module):
    r"""Simplified Mamba-2 State Space Model (SSM) block."""

    def __init__(
        self,
        d_model: int,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 2,
        **kwargs: Any,
    ) -> None:
        super().__init__()
        self.d_model: int = d_model
        self.d_state: int = d_state
        self.expand: int = expand
        inner_dim: int = d_model * expand

        self.in_proj: nn.Linear = nn.Linear(d_model, inner_dim * 2)
        self.conv1d: nn.Conv1d = nn.Conv1d(
            inner_dim, inner_dim, d_conv, groups=inner_dim, padding=d_conv - 1
        )
        self.dt_proj: nn.Linear = nn.Linear(inner_dim, inner_dim)
        self.A_log: nn.Parameter = nn.Parameter(torch.log(torch.rand(inner_dim, d_state) * 0.1))
        self.D: nn.Parameter = nn.Parameter(torch.ones(inner_dim))
        self.out_proj: nn.Linear = nn.Linear(inner_dim, d_model)

    def forward(self, x: torch.Tensor, **kwargs: Any) -> torch.Tensor:
        B, L, D = x.shape
        inner_dim = self.in_proj.out_features // 2
        x_proj = self.in_proj(x)
        x_ssm, z = x_proj.chunk(2, dim=-1)
        x_ssm = (
            self.conv1d(x_ssm.transpose(1, 2).contiguous())[..., :L].transpose(1, 2).contiguous()
        )
        dt = F.softplus(self.dt_proj(x_ssm))
        A = -torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0)).to(dtype=x.dtype)

        dA_exp_arg = torch.clamp(torch.einsum("bli,in->blin", dt, A), min=-20.0, max=0.0)
        dA = torch.exp(dA_exp_arg)
        dBX = torch.einsum("bli,blin->blin", x_ssm, dt.unsqueeze(-1))

        # Numerically stable recurrence bounded without exponential overflow
        h_t = torch.zeros(B, inner_dim, self.d_state, device=x.device, dtype=x.dtype)
        h_list = []
        for t in range(L):
            h_t = dA[:, t] * h_t + dBX[:, t]
            h_t = torch.clamp(h_t, min=-1e4, max=1e4)
            h_list.append(h_t)
        h = torch.stack(h_list, dim=1)

        out = (h * self.D.view(1, 1, inner_dim, 1)).sum(dim=-1) * F.silu(z)
        return self.out_proj(out)

Simplified Mamba-2 State Space Model (SSM) block.

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

Ancestors

• torch.nn.modules.module.Module

Methods

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

Expand source code

def forward(self, x: torch.Tensor, **kwargs: Any) -> torch.Tensor:
    B, L, D = x.shape
    inner_dim = self.in_proj.out_features // 2
    x_proj = self.in_proj(x)
    x_ssm, z = x_proj.chunk(2, dim=-1)
    x_ssm = (
        self.conv1d(x_ssm.transpose(1, 2).contiguous())[..., :L].transpose(1, 2).contiguous()
    )
    dt = F.softplus(self.dt_proj(x_ssm))
    A = -torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0)).to(dtype=x.dtype)

    dA_exp_arg = torch.clamp(torch.einsum("bli,in->blin", dt, A), min=-20.0, max=0.0)
    dA = torch.exp(dA_exp_arg)
    dBX = torch.einsum("bli,blin->blin", x_ssm, dt.unsqueeze(-1))

    # Numerically stable recurrence bounded without exponential overflow
    h_t = torch.zeros(B, inner_dim, self.d_state, device=x.device, dtype=x.dtype)
    h_list = []
    for t in range(L):
        h_t = dA[:, t] * h_t + dBX[:, t]
        h_t = torch.clamp(h_t, min=-1e4, max=1e4)
        h_list.append(h_t)
    h = torch.stack(h_list, dim=1)

    out = (h * self.D.view(1, 1, inner_dim, 1)).sum(dim=-1) * F.silu(z)
    return self.out_proj(out)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class MambaBlock (config: Any | None = None,
d_model: int | None = None,
d_state: int = 16,
d_conv: int = 4,
expand: int = 2,
dt_rank: int | None = None,
conv_bias: bool = True,
bias: bool = False,
bidirectional: bool = False,
is_decoder_override: bool | None = None,
freeze_base: bool = False,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("mamba")
@ATTENTION_REGISTRY.register("mamba_attention")
@ATTENTION_REGISTRY.register("mamba_block")
class MambaBlock(nn.Module):
    """Faithful S6 / Mamba block, production-hardened."""

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

    def __init__(
        self,
        config: Any | None = None,
        d_model: int | None = None,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 2,
        dt_rank: int | None = None,
        conv_bias: bool = True,
        bias: bool = False,
        bidirectional: bool = False,
        is_decoder_override: bool | None = None,
        freeze_base: bool = False,
        **kwargs: Any,
    ) -> None:
        # Convenience: MambaBlock(768) support
        if isinstance(config, (int, float)) and d_model is None:
            d_model = int(config)
            config = None

        super().__init__()
        self.config = config

        # ---- d_model inference (MB-11) ----
        if d_model is None:
            d_model = _infer_d_model(config, kwargs.get("source_attn"), None, kwargs)
        self.d_model = int(d_model)

        # ---- Dimensions ----
        self.d_state = int(d_state)
        self.d_conv = int(d_conv)
        self.expand = int(expand)
        self.d_inner = self.expand * self.d_model
        self.dt_rank = int(dt_rank or math.ceil(self.d_model / 16))
        self.bidirectional = bool(bidirectional)
        self.freeze_base = bool(freeze_base)

        # ---- Decoder detection (MB-10) ----
        if is_decoder_override is not None:
            self.is_decoder = bool(is_decoder_override)
        else:
            _mt = str(getattr(config, "model_type", "") or "").lower() if config else ""
            _dec_prefixes = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
            )
            self.is_decoder = (
                bool(getattr(config, "is_decoder", False) or getattr(config, "causal", False))
                if config
                else False
            ) or any(_mt.startswith(p) for p in _dec_prefixes)

        # ---- Projections ----
        self.in_proj = nn.Linear(self.d_model, 2 * self.d_inner, bias=bias)
        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            kernel_size=self.d_conv,
            groups=self.d_inner,
            padding=self.d_conv - 1,
            bias=conv_bias,
        )
        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + 2 * self.d_state, bias=False)
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias)

        # ---- S4D-real initialization of A ----
        A = torch.arange(1, self.d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))
        self.D = nn.Parameter(torch.ones(self.d_inner))

        # ---- dt initialization ----
        dt_min, dt_max = 0.001, 0.1
        dt_init = torch.exp(
            torch.rand(self.d_inner) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)
        )
        dt_init_inv = dt_init + torch.log(-torch.expm1(-dt_init))
        with torch.no_grad():
            self.dt_proj.bias.copy_(dt_init_inv)
            bound = 1.0 / math.sqrt(self.dt_rank)
            nn.init.uniform_(self.dt_proj.weight, -bound, bound)
            nn.init.kaiming_uniform_(self.conv1d.weight, a=math.sqrt(5))
            if self.conv1d.bias is not None:
                nn.init.zeros_(self.conv1d.bias)

        # ---- Bidirectional (backward SSM) ----
        if self.bidirectional:
            self.A_log_bwd = nn.Parameter(torch.log(A.clone() + 0.01 * torch.randn_like(A).abs()))
            self.dt_proj_bwd = nn.Linear(self.dt_rank, self.d_inner, bias=True)
            with torch.no_grad():
                self.dt_proj_bwd.bias.copy_(dt_init_inv)
                nn.init.uniform_(self.dt_proj_bwd.weight, -bound, bound)
            self.x_proj_bwd = nn.Linear(self.d_inner, self.dt_rank + 2 * self.d_state, bias=False)
            self.D_bwd = nn.Parameter(torch.ones(self.d_inner))

        # ---- Output-format detection (MB-12) ----
        source_attn = kwargs.get("source_attn") or kwargs.get("old_attn")
        self.is_nested_encoder = False
        self._return_format = _detect_return_format(source_attn, False)

        if freeze_base:
            self.freeze_base_weights()

    # ------------------------------------------------------------------
    def freeze_base_weights(self) -> None:
        """MB-04: match the framework's LoRA-freeze convention."""
        for p in self.parameters():
            p.requires_grad = False

    def unfreeze_base_weights(self) -> None:
        for p in self.parameters():
            p.requires_grad = True

    # ------------------------------------------------------------------
    # Weight distillation (MB-07, MB-08)
    # ------------------------------------------------------------------
    def copy_weights(self, src: nn.Module) -> None:
        """
        Distill attention weights (W_q, W_k, W_v, W_o) into Mamba's S6
        parameters via low-rank mapping:
            W_v  → in_proj (SSM input path)
            W_k  → x_proj (B branch, via SVD)
            W_q  → x_proj (C branch, via SVD)
            W_o  → out_proj
            dt   → kept at initialization (learnable, input-dependent via x_proj)
        """
        if src is None:
            return

        extracted: dict[str, torch.Tensor | None] = {
            "q": None,
            "k": None,
            "v": None,
            "o": None,
        }

        # --- 1. UniversalAdapter ---
        try:
            from neuralops.core.adapters import UniversalAdapter

            for proj in ("q", "k", "v", "o"):
                try:
                    w, _ = UniversalAdapter.get_weight_bias(src, proj)
                    if w is not None:
                        extracted[proj] = w
                except Exception:
                    pass
        except Exception:
            pass

        # --- 2. GPT-2 fused c_attn / c_proj ---
        if hasattr(src, "c_attn") and src.c_attn is not None and extracted["q"] is None:
            c = src.c_attn
            base = getattr(c, "base_layer", c)
            w = getattr(base, "weight", getattr(c, "weight", None))
            if w is not None:
                w = w.data.clone()
                is_conv = (
                    c.__class__.__name__ == "Conv1D"
                    or base.__class__.__name__ == "Conv1D"
                    or (w.ndim == 2 and w.shape[1] == 3 * self.d_model)
                )
                if is_conv:
                    w = w.t().contiguous()
                split = w.size(0) // 3
                extracted["q"] = w[:split, :]
                extracted["k"] = w[split : 2 * split, :]
                extracted["v"] = w[2 * split :, :]

        if hasattr(src, "c_proj") and src.c_proj is not None and extracted["o"] is None:
            o = src.c_proj
            base = getattr(o, "base_layer", o)
            w = getattr(base, "weight", getattr(o, "weight", None))
            if w is not None:
                w = w.data.clone()
                if w.ndim == 2 and w.shape[0] < w.shape[1]:
                    w = w.t().contiguous()
                extracted["o"] = w

        # --- 3. Flat / nested (BERT / RoBERTa) ---
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        nested = getattr(src, "self", None)
        for proj, names in named_map.items():
            if extracted[proj] is not None:
                continue
            # Try nested first
            if nested is not None:
                for n in names:
                    mod = getattr(nested, n, None)
                    if mod is not None and hasattr(mod, "weight"):
                        extracted[proj] = mod.weight.data
                        break
            if extracted[proj] is not None:
                continue
            # Then flat on src
            for n in names:
                mod = getattr(src, n, None)
                if mod is not None and hasattr(mod, "weight"):
                    extracted[proj] = mod.weight.data
                    break
            # Nested output.dense for "o"
            if extracted["o"] is None and proj == "o":
                out = getattr(src, "output", None)
                dense = getattr(out, "dense", None) if out is not None else None
                if dense is not None and hasattr(dense, "weight"):
                    extracted["o"] = dense.weight.data

        # Mark nested encoder for later (MB-05)
        cls_name = src.__class__.__name__
        if cls_name in (
            "BertSelfAttention",
            "RobertaSelfAttention",
            "XLMRobertaSelfAttention",
        ) or (nested is not None and hasattr(nested, "query")):
            self.is_nested_encoder = True

        # ---- Apply mapping (MB-07, MB-08) ----
        dtype = self.in_proj.weight.dtype
        device = self.in_proj.weight.device

        def _to_t(x: torch.Tensor | None, r: int, c: int) -> torch.Tensor | None:
            """Project W (any shape) to target shape (r, c) via SVD."""
            if x is None:
                return None
            x = x.float()
            if x.ndim != 2:
                return None
            M, N = x.shape
            u, s, vh = torch.linalg.svd(x, full_matrices=False)
            k = min(r, c, s.numel())

            u_r = u[: min(r, M), :k]
            vh_c = vh[:k, : min(c, N)]
            core = (u_r * s[:k].unsqueeze(0)) @ vh_c

            out = torch.zeros(r, c, dtype=torch.float32, device=x.device)
            out[: core.size(0), : core.size(1)] = core

            if r > core.size(0):
                rem_r = r - core.size(0)
                reps_r = (rem_r + core.size(0) - 1) // core.size(0)
                tiled_r = core.repeat(reps_r, 1)[:rem_r, :]
                out[core.size(0) :, : core.size(1)] = tiled_r

            if c > core.size(1):
                rem_c = c - core.size(1)
                reps_c = (rem_c + core.size(1) - 1) // core.size(1)
                tiled_c = out[:, : core.size(1)].repeat(1, reps_c)[:, :rem_c]
                out[:, core.size(1) :] = tiled_c

            return out.to(device=device, dtype=dtype)

        with torch.no_grad():
            w_q = extracted["q"]
            w_k = extracted["k"]
            w_v = extracted["v"]
            w_o = extracted["o"]

            # ---------- in_proj: (2*d_inner, d_model) ----------
            # x_ssm path = rows [0 : d_inner]
            # z     path = rows [d_inner : 2*d_inner]
            if w_v is not None:
                w_v_p = _to_t(w_v, self.d_inner, self.d_model)
                if w_v_p is not None:
                    # x_ssm: full projection from W_v
                    self.in_proj.weight.data[: self.d_inner, : self.d_model].copy_(w_v_p)
                    # z: scaled copy so the gate starts balanced
                    self.in_proj.weight.data[self.d_inner :, : self.d_model].copy_(0.5 * w_v_p)

            # ---------- x_proj: (dt_rank + 2*d_state, d_inner) ----------
            # rows [0 : dt_rank]                     → dt  (keep init)
            # rows [dt_rank : dt_rank + d_state]     → B   (from W_k)
            # rows [dt_rank + d_state : ...]         → C   (from W_q)
            # MB-01: DO NOT zero the whole x_proj; only reset B/C rows.
            with torch.no_grad():
                b_start = self.dt_rank
                b_end = self.dt_rank + self.d_state
                c_start = b_end
                c_end = c_start + self.d_state

                if w_k is not None:
                    b_mat = _to_t(w_k, self.d_state, self.d_inner)
                    if b_mat is not None:
                        self.x_proj.weight.data[b_start:b_end, :].copy_(b_mat)
                if w_q is not None:
                    c_mat = _to_t(w_q, self.d_state, self.d_inner)
                    if c_mat is not None:
                        self.x_proj.weight.data[c_start:c_end, :].copy_(c_mat)

            # ---------- out_proj: (d_model, d_inner) ----------
            if w_o is not None:
                w_o_p = _to_t(w_o, self.d_model, self.d_inner)
                if w_o_p is not None:
                    self.out_proj.weight.data.copy_(w_o_p)
                    if self.out_proj.bias is not None:
                        self.out_proj.bias.data.zero_()

            # ---------- conv1d: identity-at-init passthrough ----------
            self.conv1d.weight.data.zero_()
            if self.conv1d.weight.dim() == 3:
                # shape (d_inner, 1, d_conv)
                self.conv1d.weight.data[:, 0, -1] = 1.0
            if self.conv1d.bias is not None:
                self.conv1d.bias.data.zero_()

        # ---- Sanity check ----
        for name, mod in (
            ("in_proj", self.in_proj),
            ("x_proj", self.x_proj),
            ("out_proj", self.out_proj),
        ):
            n = float(torch.linalg.norm(mod.weight.data.float()).item())
            if n < 1e-6:
                logger.warning("MambaBlock: '%s' near-zero norm (%.2e)", name, n)

        if self.freeze_base:
            self.freeze_base_weights()

    # ------------------------------------------------------------------
    # State helpers
    # ------------------------------------------------------------------
    def _read_state(
        self, past_state: Any
    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, int]:
        """Returns (h, conv_state, dt_state, past_len)."""
        if past_state is None:
            return None, None, None, 0
        if isinstance(past_state, (tuple, list)):
            if len(past_state) >= 3:
                h = past_state[0]
                conv_state = past_state[1]
                dt_state = past_state[2] if len(past_state) > 2 else None
                past_len = h.shape[-1] if h is not None else 0
                # past_len from conv_state if available
                if conv_state is not None:
                    past_len = conv_state.shape[-1] - (self.d_conv - 1)
                    past_len = max(0, past_len)
                return h, conv_state, dt_state, past_len
            elif len(past_state) == 2:
                return past_state[0], past_state[1], None, 0
            elif len(past_state) == 1:
                return past_state[0], None, None, 0
        return None, None, None, 0

    def _recurrent_scan(
        self,
        u: torch.Tensor,  # (B, d_inner, L)
        delta: torch.Tensor,  # (B, d_inner, L)
        A: torch.Tensor,  # (d_inner, d_state)
        B_ssm: torch.Tensor,  # (B, d_state, L)
        C_ssm: torch.Tensor,  # (B, d_state, L)
        D: torch.Tensor,  # (d_inner,)
        state: torch.Tensor | None,  # (B, d_inner, d_state)
        chunk_size: int = 32,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Recurrent selective scan with chunked activation checkpointing and bmm.

        Fix 1: Split L into chunks of size chunk_size and checkpoint each chunk
               during training to bound autograd memory to O(chunk_size).
        Fix 2: Use torch.bmm instead of torch.einsum inside the loop for cuBLAS
               acceleration and shared autograd backward kernels.
        """
        B, d_inner, L = u.shape
        d_state = A.shape[1]

        target_dtype = u.dtype
        h = (
            state.to(target_dtype)
            if state is not None
            else torch.zeros(B, d_inner, d_state, device=u.device, dtype=target_dtype)
        )

        delta = F.softplus(delta).to(target_dtype)
        delta_a = delta.unsqueeze(-1) * A.view(1, d_inner, 1, d_state).to(target_dtype)
        delta_a = torch.exp(delta_a).to(target_dtype)
        delta_b_u = (
            delta.unsqueeze(-1)
            * B_ssm.permute(0, 2, 1).unsqueeze(1).to(target_dtype)
            * u.unsqueeze(-1)
        ).to(target_dtype)
        C_in = C_ssm.to(target_dtype)

        def _scan_chunk(
            h_init: torch.Tensor,
            da_c: torch.Tensor,
            dbu_c: torch.Tensor,
            c_c: torch.Tensor,
        ) -> tuple[torch.Tensor, torch.Tensor]:
            L_c = da_c.shape[2]
            ys_c = []
            curr_h = h_init
            for i in range(L_c):
                curr_h = da_c[:, :, i, :] * curr_h + dbu_c[:, :, i, :]
                c_t = (
                    c_c[:, :, i].contiguous().unsqueeze(-1)
                )  # (B, d_state, 1) contiguous for cuBLAS
                y_t = torch.bmm(curr_h, c_t).squeeze(-1)  # (B, d_inner)
                ys_c.append(y_t)
            return torch.stack(ys_c, dim=-1), curr_h

        # Chunked scan with activation checkpointing during training
        ys_all = []
        import torch.utils.checkpoint as cp

        use_cp = self.training and torch.is_grad_enabled()

        for s in range(0, L, chunk_size):
            e = min(s + chunk_size, L)
            da_slice = delta_a[:, :, s:e, :]
            dbu_slice = delta_b_u[:, :, s:e, :]
            c_slice = C_in[:, :, s:e]

            if use_cp:
                y_chunk, h = cp.checkpoint(
                    _scan_chunk,
                    h,
                    da_slice,
                    dbu_slice,
                    c_slice,
                    use_reentrant=False,
                )
            else:
                y_chunk, h = _scan_chunk(
                    h,
                    da_slice,
                    dbu_slice,
                    c_slice,
                )
            ys_all.append(y_chunk)

        y = torch.cat(ys_all, dim=-1)  # (B, d_inner, L)
        y = y + u * D.view(1, d_inner, 1).to(target_dtype)
        return y, h

    def _parallel_scan(
        self,
        u: torch.Tensor,
        delta: torch.Tensor,
        A: torch.Tensor,
        B_ssm: torch.Tensor,
        C_ssm: torch.Tensor,
        D: torch.Tensor,
    ) -> torch.Tensor:
        """
        Parallel scan dispatch with 3-tier fallback:
          Path 1: Official CUDA kernel via selective_scan (Linux only, when compiled)
          Path 2: Triton selective_scan (if mamba_ssm Triton ops available)
          Path 3: Fast chunked parallel scan (framework.core.selective_scan, O(1) ops per chunk)
          Path 4: Chunked recurrent scan with activation checkpointing (failsafe fallback)
        """
        # Path 1: Official CUDA kernel
        if HAS_MAMBA_CUDA and selective_scan is not None:
            try:
                return selective_scan(
                    u=u,
                    delta=delta,
                    A=A,
                    B=B_ssm,
                    C=C_ssm,
                    D=D,
                    delta_softplus=True,
                )
            except Exception as e:
                logger.debug("CUDA selective_scan failed: %s", e)

        # Path 2: Triton kernel
        if HAS_TRITON_MAMBA and triton_selective_scan is not None:
            try:
                y = triton_selective_scan(
                    u.transpose(1, 2).contiguous(),
                    delta.transpose(1, 2).contiguous(),
                    A,
                    B_ssm.transpose(1, 2).contiguous(),
                    C_ssm.transpose(1, 2).contiguous(),
                    D,
                    delta_softplus=True,
                )
                return y.transpose(1, 2).contiguous()
            except Exception as e:
                logger.warning("Triton selective_scan failed: %s", e)

        # Path 3: Fast chunked parallel scan (avoids token-by-token Python loop)
        if selective_scan is not None:
            try:
                return selective_scan(
                    u=u,
                    delta=delta,
                    A=A,
                    B=B_ssm,
                    C=C_ssm,
                    D=D,
                    delta_softplus=True,
                )
            except Exception as e:
                logger.debug("Parallel selective_scan fallback failed: %s", e)

        # Path 4: Failsafe sequential / chunked recurrent fallback
        y, _ = self._recurrent_scan(u, delta, A, B_ssm, C_ssm, D, None)
        return y

    # ------------------------------------------------------------------
    # Forward (MB-02, MB-03, MB-09, MB-12)
    # ------------------------------------------------------------------
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        *args: Any,
        past_state: Any | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool = False,
        **kwargs: Any,
    ):
        if kwargs.get("encoder_hidden_states") is not None:
            raise NotImplementedError("MambaBlock does not support cross-attention")

        # Backward-compat: accept past_key_value as an alias for past_state
        if past_state is None and past_key_value is not None:
            past_state = past_key_value

        B, L, _ = hidden_states.shape
        device = hidden_states.device
        target_dtype = hidden_states.dtype

        # ---- Read state ----
        h_prev, conv_prev, _, past_len = self._read_state(past_state)

        # ---- 1. Input projection ----
        xz = self.in_proj(hidden_states)  # (B, L, 2*d_inner)
        x_ssm, z = xz.chunk(2, dim=-1)  # each (B, L, d_inner)

        # ---- 2. Depthwise conv1d with streaming state ----
        if conv_prev is not None:
            # conv_prev: (B, d_inner, d_conv - 1)
            # Concatenate along time axis
            x_cat = torch.cat(
                [conv_prev.transpose(1, 2), x_ssm], dim=1
            )  # (B, L + d_conv - 1, d_inner)
        else:
            x_cat = x_ssm

        x_conv_full = self.conv1d(x_cat.transpose(1, 2))  # (B, d_inner, L' + d_conv - 1)
        # Take only the L last positions
        x_conv = x_conv_full[..., -L:].transpose(1, 2)  # (B, L, d_inner)

        # New conv state: last (d_conv - 1) tokens of input
        if self.d_conv - 1 <= L:
            conv_next = x_cat[:, -(self.d_conv - 1) :, :].transpose(1, 2).contiguous()
        else:
            # pad with zero if sequence shorter than kernel
            pad = self.d_conv - 1 - L
            conv_next = (
                torch.cat(
                    [
                        torch.zeros(B, pad, self.d_inner, device=device, dtype=x_cat.dtype),
                        x_cat,
                    ],
                    dim=1,
                )
                .transpose(1, 2)
                .contiguous()
            )

        x_ssm = F.silu(x_conv)

        # ---- 3. Selective parameters ----
        x_dbl = self.x_proj(x_ssm)  # (B, L, dt_rank + 2*d_state)
        dt, B_ssm, C_ssm = x_dbl.split([self.dt_rank, self.d_state, self.d_state], dim=-1)
        dt = self.dt_proj(dt)  # (B, L, d_inner)

        A = -torch.exp(self.A_log)  # (d_inner, d_state)

        # ---- 4. Selective scan ----
        u_t = x_ssm.transpose(1, 2)  # (B, d_inner, L)
        dt_t = dt.transpose(1, 2)
        B_t = B_ssm.transpose(1, 2)
        C_t = C_ssm.transpose(1, 2)

        if h_prev is None:
            y = self._parallel_scan(u_t, dt_t, A, B_t, C_t, self.D)
            h_next = None
        else:
            y, h_next = self._recurrent_scan(u_t, dt_t, A, B_t, C_t, self.D, h_prev)

        y = y.transpose(1, 2)  # (B, L, d_inner)

        # ---- 5. Bidirectional (encoder only) ----
        if self.bidirectional and not self.is_decoder:
            x_rev = torch.flip(x_ssm, dims=[1])
            x_dbl_rev = self.x_proj_bwd(x_rev)
            dt_b, B_b, C_b = x_dbl_rev.split([self.dt_rank, self.d_state, self.d_state], dim=-1)
            dt_b = self.dt_proj_bwd(dt_b)
            A_b = -torch.exp(self.A_log_bwd)
            y_b = self._parallel_scan(
                torch.flip(x_ssm, dims=[1]).transpose(1, 2),
                dt_b.transpose(1, 2),
                A_b,
                B_b.transpose(1, 2),
                C_b.transpose(1, 2),
                self.D_bwd,
            )
            y_b = torch.flip(y_b.transpose(1, 2), dims=[1])
            y = 0.5 * (y + y_b)

        # ---- 6. Gating ----
        y = y.to(target_dtype) * F.silu(z.to(target_dtype))

        # ---- 7. Output projection ----
        if self.is_nested_encoder:
            # Nested encoder: outer layer will project. Skip out_proj.
            out = y.to(target_dtype)
        else:
            out = self.out_proj(y).to(target_dtype)

        # ---- 8. attention_mask (MB-03, MB-09) ----
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                out = out * attention_mask.unsqueeze(-1).to(out.dtype)
            elif attention_mask.dim() == 3:
                # (B, Lq, Lk) — take the last column (query's validity) or diagonal
                m = (
                    attention_mask.diagonal(dim1=-2, dim2=-1)
                    if attention_mask.size(-1) == L
                    else attention_mask[..., -1]
                )
                out = out * m.unsqueeze(-1).to(out.dtype)
            elif attention_mask.dim() == 4:
                # (B, 1, Lq, Lk) — same reduction
                m = attention_mask.squeeze(1)
                m = m.diagonal(dim1=-2, dim2=-1) if m.size(-1) == L else m[..., -1]
                out = out * m.unsqueeze(-1).to(out.dtype)

        # ---- 9. Output tuple (MB-12) ----
        present_state = (h_next, conv_next)
        attn_weights = None

        fmt = self._return_format
        if fmt in ("bert", "gpt2"):
            return (out, attn_weights)
        # llama
        return (out, attn_weights, present_state)

Faithful S6 / Mamba block, production-hardened.

Ancestors

• 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.

Methods

def copy_weights(self, src: nn.Module) ‑> None

Expand source code

def copy_weights(self, src: nn.Module) -> None:
    """
    Distill attention weights (W_q, W_k, W_v, W_o) into Mamba's S6
    parameters via low-rank mapping:
        W_v  → in_proj (SSM input path)
        W_k  → x_proj (B branch, via SVD)
        W_q  → x_proj (C branch, via SVD)
        W_o  → out_proj
        dt   → kept at initialization (learnable, input-dependent via x_proj)
    """
    if src is None:
        return

    extracted: dict[str, torch.Tensor | None] = {
        "q": None,
        "k": None,
        "v": None,
        "o": None,
    }

    # --- 1. UniversalAdapter ---
    try:
        from neuralops.core.adapters import UniversalAdapter

        for proj in ("q", "k", "v", "o"):
            try:
                w, _ = UniversalAdapter.get_weight_bias(src, proj)
                if w is not None:
                    extracted[proj] = w
            except Exception:
                pass
    except Exception:
        pass

    # --- 2. GPT-2 fused c_attn / c_proj ---
    if hasattr(src, "c_attn") and src.c_attn is not None and extracted["q"] is None:
        c = src.c_attn
        base = getattr(c, "base_layer", c)
        w = getattr(base, "weight", getattr(c, "weight", None))
        if w is not None:
            w = w.data.clone()
            is_conv = (
                c.__class__.__name__ == "Conv1D"
                or base.__class__.__name__ == "Conv1D"
                or (w.ndim == 2 and w.shape[1] == 3 * self.d_model)
            )
            if is_conv:
                w = w.t().contiguous()
            split = w.size(0) // 3
            extracted["q"] = w[:split, :]
            extracted["k"] = w[split : 2 * split, :]
            extracted["v"] = w[2 * split :, :]

    if hasattr(src, "c_proj") and src.c_proj is not None and extracted["o"] is None:
        o = src.c_proj
        base = getattr(o, "base_layer", o)
        w = getattr(base, "weight", getattr(o, "weight", None))
        if w is not None:
            w = w.data.clone()
            if w.ndim == 2 and w.shape[0] < w.shape[1]:
                w = w.t().contiguous()
            extracted["o"] = w

    # --- 3. Flat / nested (BERT / RoBERTa) ---
    named_map = {
        "q": ("q_proj", "q_lin", "query"),
        "k": ("k_proj", "k_lin", "key"),
        "v": ("v_proj", "v_lin", "value"),
        "o": ("o_proj", "out_proj", "out_lin", "dense"),
    }
    nested = getattr(src, "self", None)
    for proj, names in named_map.items():
        if extracted[proj] is not None:
            continue
        # Try nested first
        if nested is not None:
            for n in names:
                mod = getattr(nested, n, None)
                if mod is not None and hasattr(mod, "weight"):
                    extracted[proj] = mod.weight.data
                    break
        if extracted[proj] is not None:
            continue
        # Then flat on src
        for n in names:
            mod = getattr(src, n, None)
            if mod is not None and hasattr(mod, "weight"):
                extracted[proj] = mod.weight.data
                break
        # Nested output.dense for "o"
        if extracted["o"] is None and proj == "o":
            out = getattr(src, "output", None)
            dense = getattr(out, "dense", None) if out is not None else None
            if dense is not None and hasattr(dense, "weight"):
                extracted["o"] = dense.weight.data

    # Mark nested encoder for later (MB-05)
    cls_name = src.__class__.__name__
    if cls_name in (
        "BertSelfAttention",
        "RobertaSelfAttention",
        "XLMRobertaSelfAttention",
    ) or (nested is not None and hasattr(nested, "query")):
        self.is_nested_encoder = True

    # ---- Apply mapping (MB-07, MB-08) ----
    dtype = self.in_proj.weight.dtype
    device = self.in_proj.weight.device

    def _to_t(x: torch.Tensor | None, r: int, c: int) -> torch.Tensor | None:
        """Project W (any shape) to target shape (r, c) via SVD."""
        if x is None:
            return None
        x = x.float()
        if x.ndim != 2:
            return None
        M, N = x.shape
        u, s, vh = torch.linalg.svd(x, full_matrices=False)
        k = min(r, c, s.numel())

        u_r = u[: min(r, M), :k]
        vh_c = vh[:k, : min(c, N)]
        core = (u_r * s[:k].unsqueeze(0)) @ vh_c

        out = torch.zeros(r, c, dtype=torch.float32, device=x.device)
        out[: core.size(0), : core.size(1)] = core

        if r > core.size(0):
            rem_r = r - core.size(0)
            reps_r = (rem_r + core.size(0) - 1) // core.size(0)
            tiled_r = core.repeat(reps_r, 1)[:rem_r, :]
            out[core.size(0) :, : core.size(1)] = tiled_r

        if c > core.size(1):
            rem_c = c - core.size(1)
            reps_c = (rem_c + core.size(1) - 1) // core.size(1)
            tiled_c = out[:, : core.size(1)].repeat(1, reps_c)[:, :rem_c]
            out[:, core.size(1) :] = tiled_c

        return out.to(device=device, dtype=dtype)

    with torch.no_grad():
        w_q = extracted["q"]
        w_k = extracted["k"]
        w_v = extracted["v"]
        w_o = extracted["o"]

        # ---------- in_proj: (2*d_inner, d_model) ----------
        # x_ssm path = rows [0 : d_inner]
        # z     path = rows [d_inner : 2*d_inner]
        if w_v is not None:
            w_v_p = _to_t(w_v, self.d_inner, self.d_model)
            if w_v_p is not None:
                # x_ssm: full projection from W_v
                self.in_proj.weight.data[: self.d_inner, : self.d_model].copy_(w_v_p)
                # z: scaled copy so the gate starts balanced
                self.in_proj.weight.data[self.d_inner :, : self.d_model].copy_(0.5 * w_v_p)

        # ---------- x_proj: (dt_rank + 2*d_state, d_inner) ----------
        # rows [0 : dt_rank]                     → dt  (keep init)
        # rows [dt_rank : dt_rank + d_state]     → B   (from W_k)
        # rows [dt_rank + d_state : ...]         → C   (from W_q)
        # MB-01: DO NOT zero the whole x_proj; only reset B/C rows.
        with torch.no_grad():
            b_start = self.dt_rank
            b_end = self.dt_rank + self.d_state
            c_start = b_end
            c_end = c_start + self.d_state

            if w_k is not None:
                b_mat = _to_t(w_k, self.d_state, self.d_inner)
                if b_mat is not None:
                    self.x_proj.weight.data[b_start:b_end, :].copy_(b_mat)
            if w_q is not None:
                c_mat = _to_t(w_q, self.d_state, self.d_inner)
                if c_mat is not None:
                    self.x_proj.weight.data[c_start:c_end, :].copy_(c_mat)

        # ---------- out_proj: (d_model, d_inner) ----------
        if w_o is not None:
            w_o_p = _to_t(w_o, self.d_model, self.d_inner)
            if w_o_p is not None:
                self.out_proj.weight.data.copy_(w_o_p)
                if self.out_proj.bias is not None:
                    self.out_proj.bias.data.zero_()

        # ---------- conv1d: identity-at-init passthrough ----------
        self.conv1d.weight.data.zero_()
        if self.conv1d.weight.dim() == 3:
            # shape (d_inner, 1, d_conv)
            self.conv1d.weight.data[:, 0, -1] = 1.0
        if self.conv1d.bias is not None:
            self.conv1d.bias.data.zero_()

    # ---- Sanity check ----
    for name, mod in (
        ("in_proj", self.in_proj),
        ("x_proj", self.x_proj),
        ("out_proj", self.out_proj),
    ):
        n = float(torch.linalg.norm(mod.weight.data.float()).item())
        if n < 1e-6:
            logger.warning("MambaBlock: '%s' near-zero norm (%.2e)", name, n)

    if self.freeze_base:
        self.freeze_base_weights()

Distill attention weights (W_q, W_k, W_v, W_o) into Mamba's S6 parameters via low-rank mapping: W_v → in_proj (SSM input path) W_k → x_proj (B branch, via SVD) W_q → x_proj (C branch, via SVD) W_o → out_proj dt → kept at initialization (learnable, input-dependent via x_proj)

def forward(self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
*args: Any,
past_state: Any | None = None,
past_key_value: Any | None = None,
output_attentions: bool = False,
**kwargs: Any) ‑> Callable[..., typing.Any]

Expand source code

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    *args: Any,
    past_state: Any | None = None,
    past_key_value: Any | None = None,
    output_attentions: bool = False,
    **kwargs: Any,
):
    if kwargs.get("encoder_hidden_states") is not None:
        raise NotImplementedError("MambaBlock does not support cross-attention")

    # Backward-compat: accept past_key_value as an alias for past_state
    if past_state is None and past_key_value is not None:
        past_state = past_key_value

    B, L, _ = hidden_states.shape
    device = hidden_states.device
    target_dtype = hidden_states.dtype

    # ---- Read state ----
    h_prev, conv_prev, _, past_len = self._read_state(past_state)

    # ---- 1. Input projection ----
    xz = self.in_proj(hidden_states)  # (B, L, 2*d_inner)
    x_ssm, z = xz.chunk(2, dim=-1)  # each (B, L, d_inner)

    # ---- 2. Depthwise conv1d with streaming state ----
    if conv_prev is not None:
        # conv_prev: (B, d_inner, d_conv - 1)
        # Concatenate along time axis
        x_cat = torch.cat(
            [conv_prev.transpose(1, 2), x_ssm], dim=1
        )  # (B, L + d_conv - 1, d_inner)
    else:
        x_cat = x_ssm

    x_conv_full = self.conv1d(x_cat.transpose(1, 2))  # (B, d_inner, L' + d_conv - 1)
    # Take only the L last positions
    x_conv = x_conv_full[..., -L:].transpose(1, 2)  # (B, L, d_inner)

    # New conv state: last (d_conv - 1) tokens of input
    if self.d_conv - 1 <= L:
        conv_next = x_cat[:, -(self.d_conv - 1) :, :].transpose(1, 2).contiguous()
    else:
        # pad with zero if sequence shorter than kernel
        pad = self.d_conv - 1 - L
        conv_next = (
            torch.cat(
                [
                    torch.zeros(B, pad, self.d_inner, device=device, dtype=x_cat.dtype),
                    x_cat,
                ],
                dim=1,
            )
            .transpose(1, 2)
            .contiguous()
        )

    x_ssm = F.silu(x_conv)

    # ---- 3. Selective parameters ----
    x_dbl = self.x_proj(x_ssm)  # (B, L, dt_rank + 2*d_state)
    dt, B_ssm, C_ssm = x_dbl.split([self.dt_rank, self.d_state, self.d_state], dim=-1)
    dt = self.dt_proj(dt)  # (B, L, d_inner)

    A = -torch.exp(self.A_log)  # (d_inner, d_state)

    # ---- 4. Selective scan ----
    u_t = x_ssm.transpose(1, 2)  # (B, d_inner, L)
    dt_t = dt.transpose(1, 2)
    B_t = B_ssm.transpose(1, 2)
    C_t = C_ssm.transpose(1, 2)

    if h_prev is None:
        y = self._parallel_scan(u_t, dt_t, A, B_t, C_t, self.D)
        h_next = None
    else:
        y, h_next = self._recurrent_scan(u_t, dt_t, A, B_t, C_t, self.D, h_prev)

    y = y.transpose(1, 2)  # (B, L, d_inner)

    # ---- 5. Bidirectional (encoder only) ----
    if self.bidirectional and not self.is_decoder:
        x_rev = torch.flip(x_ssm, dims=[1])
        x_dbl_rev = self.x_proj_bwd(x_rev)
        dt_b, B_b, C_b = x_dbl_rev.split([self.dt_rank, self.d_state, self.d_state], dim=-1)
        dt_b = self.dt_proj_bwd(dt_b)
        A_b = -torch.exp(self.A_log_bwd)
        y_b = self._parallel_scan(
            torch.flip(x_ssm, dims=[1]).transpose(1, 2),
            dt_b.transpose(1, 2),
            A_b,
            B_b.transpose(1, 2),
            C_b.transpose(1, 2),
            self.D_bwd,
        )
        y_b = torch.flip(y_b.transpose(1, 2), dims=[1])
        y = 0.5 * (y + y_b)

    # ---- 6. Gating ----
    y = y.to(target_dtype) * F.silu(z.to(target_dtype))

    # ---- 7. Output projection ----
    if self.is_nested_encoder:
        # Nested encoder: outer layer will project. Skip out_proj.
        out = y.to(target_dtype)
    else:
        out = self.out_proj(y).to(target_dtype)

    # ---- 8. attention_mask (MB-03, MB-09) ----
    if attention_mask is not None:
        if attention_mask.dim() == 2:
            out = out * attention_mask.unsqueeze(-1).to(out.dtype)
        elif attention_mask.dim() == 3:
            # (B, Lq, Lk) — take the last column (query's validity) or diagonal
            m = (
                attention_mask.diagonal(dim1=-2, dim2=-1)
                if attention_mask.size(-1) == L
                else attention_mask[..., -1]
            )
            out = out * m.unsqueeze(-1).to(out.dtype)
        elif attention_mask.dim() == 4:
            # (B, 1, Lq, Lk) — same reduction
            m = attention_mask.squeeze(1)
            m = m.diagonal(dim1=-2, dim2=-1) if m.size(-1) == L else m[..., -1]
            out = out * m.unsqueeze(-1).to(out.dtype)

    # ---- 9. Output tuple (MB-12) ----
    present_state = (h_next, conv_next)
    attn_weights = None

    fmt = self._return_format
    if fmt in ("bert", "gpt2"):
        return (out, attn_weights)
    # llama
    return (out, attn_weights, present_state)

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    """MB-04: match the framework's LoRA-freeze convention."""
    for p in self.parameters():
        p.requires_grad = False

MB-04: match the framework's LoRA-freeze convention.

def unfreeze_base_weights(self) ‑> None

Expand source code

def unfreeze_base_weights(self) -> None:
    for p in self.parameters():
        p.requires_grad = True

class MambaSVDFlash (config: Any | None = None,
rank: int | None = 16,
d_model: int | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
comp_config: Any | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
bottleneck_dim: int = 64,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("mamba_svd_flash", supports_causal=True, supports_gqa=True)
class MambaSVDFlash(BaseAttention):
    """
    Hybrid SSM + MHA with per-chunk SVD and learnable per-head gating.

    State cache layout (for streaming):
        past_key_value = {
            "ssm_h": (B, D),           SSM hidden state h_{L-1}
            "attn_k": (B, H, L, Dh),   attention keys
            "attn_v": (B, H, L, Dh),   attention values
            "seq_len": int,
        }
    Or tuple form: (ssm_h, attn_k, attn_v, seq_len).
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 16,
        d_model: int | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        comp_config: Any | None = None,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        bottleneck_dim: int = 64,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn

        # ---- d_model ----
        dim = (
            kwargs.get("d_model")
            or kwargs.get("hidden_size")
            or kwargs.get("dim")
            or d_model
            or hidden_size
        )
        if dim is None:
            if config is not None:
                if isinstance(config, int):
                    dim = config
                elif isinstance(config, dict):
                    dim = config.get("hidden_size", config.get("dim", 128))
                else:
                    dim = getattr(config, "hidden_size", getattr(config, "dim", 128))
            elif source_attn is not None:
                cand = getattr(
                    source_attn,
                    "hidden_size",
                    getattr(
                        source_attn,
                        "embed_dim",
                        getattr(source_attn, "all_head_size", None),
                    ),
                )
                if cand is None and hasattr(source_attn, "q_proj"):
                    q = getattr(source_attn, "q_proj", None)
                    if q is not None and hasattr(q, "in_features"):
                        cand = q.in_features
                dim = cand or 128

        config_obj = None if isinstance(config, int) else config
        super().__init__(config_obj, rank if rank is not None else 16)

        self.d_model = int(dim or 128)
        self.dim = self.d_model
        self.hidden_size = self.d_model
        self.rank = max(1, int(rank if rank is not None else 16))  # MSF-FIX-02

        # ---- LoRA config ----
        if config_obj is not None:
            self.lora_rank = int(
                getattr(
                    config_obj,
                    "lora_rank",
                    getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
                )
            )
            self.lora_alpha = float(
                getattr(
                    config_obj,
                    "lora_alpha",
                    getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
                )
            )
        elif comp_config is not None:
            self.lora_rank = int(getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)))
            self.lora_alpha = float(
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0))
            )
        else:
            self.lora_rank = int(kwargs.get("lora_rank", 16))
            self.lora_alpha = float(kwargs.get("lora_alpha", 32.0))
        self.freeze_base = self.lora_rank > 0

        # ---- decoder detection ----
        model_type = (
            str(getattr(config_obj, "model_type", "") or "").lower()
            if config_obj is not None
            else ""
        )
        is_dec = False
        if config_obj is not None:
            is_dec = bool(
                getattr(config_obj, "is_decoder", False) or getattr(config_obj, "is_causal", False)
            )
        if any(
            model_type.startswith(p)
            for p in (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        if source_attn is not None:
            if hasattr(source_attn, "c_attn") or getattr(source_attn, "is_decoder", False):
                is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # ---- heads / GQA ----
        n_heads = num_heads or kwargs.get("n_heads") or kwargs.get("num_attention_heads")
        if n_heads is None:
            if config_obj is not None:
                n_heads = getattr(
                    config_obj,
                    "num_attention_heads",
                    getattr(config_obj, "n_heads", 12),
                )
            elif source_attn is not None:
                n_heads = getattr(
                    source_attn,
                    "num_attention_heads",
                    getattr(source_attn, "num_heads", 12),
                )
            else:
                n_heads = 12
        self.num_heads = int(n_heads)
        if self.d_model % self.num_heads != 0:
            logger.warning("MambaSVDFlash: d_model %% num_heads != 0; forcing num_heads=1")
            self.num_heads = 1
        self.head_dim = self.d_model // max(1, self.num_heads)

        cfg_kv = None
        if config_obj is not None:
            cfg_kv = getattr(
                config_obj,
                "num_key_value_heads",
                getattr(config_obj, "num_kv_heads", None),
            )
        if cfg_kv is None:
            cfg_kv = kwargs.get("num_kv_heads", self.num_heads)
        self.num_kv_heads = max(1, min(int(cfg_kv), self.num_heads))
        self.kv_repeat = self.num_heads // self.num_kv_heads
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        # ---- SSM sub-module (MSF-FIX-10) ----
        bd = int(
            getattr(comp_config, "bottleneck_dim", bottleneck_dim)
            if comp_config
            else bottleneck_dim
        )
        self.ssm_fwd = SVDS6StateSpace(self.d_model, rank=self.rank, bottleneck_dim=bd)

        if not self.is_decoder:
            self.ssm_bwd = SVDS6StateSpace(self.d_model, rank=self.rank, bottleneck_dim=bd)
            # MSF-FIX-13: per-head mix
            self.bidirectional_mix = nn.Parameter(torch.zeros(self.num_heads))
        else:
            self.ssm_bwd = None
            self.bidirectional_mix = None

        # ---- Attention projections ----
        self.q_proj = LoRALinear(
            self.d_model,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.d_model,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.d_model,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.d_model,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # ---- Gating (MSF-FIX-12: LoRALinear, per-head output) ----
        self.gating = LoRALinear(
            self.d_model,
            self.num_heads,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        with torch.no_grad():
            nn.init.zeros_(self.gating.weight)
            if self.gating.bias is not None:
                nn.init.constant_(self.gating.bias, -2.0)

        # ---- Dropout ----
        dropout_p = 0.1
        if config_obj is not None:
            dropout_p = float(
                getattr(
                    config_obj,
                    "dropout",
                    getattr(config_obj, "attention_probs_dropout_prob", 0.1),
                )
            )
        elif comp_config is not None:
            dropout_p = float(getattr(comp_config, "dropout", 0.1))
        else:
            dropout_p = float(kwargs.get("dropout", 0.1))
        self.dropout = nn.Dropout(dropout_p)

        self.source_attn = source_attn
        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )

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

    # ---- aliases ----
    @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 query(self):
        return self.q_proj

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

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

    # ------------------------------------------------------------------
    def freeze_base_weights(self):
        # MSF-FIX-09: freeze SSM projections too
        for m in (
            self.q_proj,
            self.k_proj,
            self.v_proj,
            self.out_proj,
            self.gating,
            self.ssm_fwd.down,
            self.ssm_fwd.up_u,
            self.ssm_fwd.up_v,
            self.ssm_fwd.up_sigma,
            self.ssm_fwd.up_b,
            self.ssm_fwd.up_c,
        ):
            for p in m.parameters():
                p.requires_grad = False
        if self.ssm_bwd is not None:
            for p in self.ssm_bwd.parameters():
                p.requires_grad = False
        # Keep LoRA adapters trainable
        for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj, self.gating):
            for name, p in m.named_parameters():
                if "lora_" in name:
                    p.requires_grad = True

    def unfreeze_base_weights(self):
        for m in (
            self.q_proj,
            self.k_proj,
            self.v_proj,
            self.out_proj,
            self.gating,
            self.ssm_fwd.down,
            self.ssm_fwd.up_u,
            self.ssm_fwd.up_v,
            self.ssm_fwd.up_sigma,
            self.ssm_fwd.up_b,
            self.ssm_fwd.up_c,
        ):
            for p in m.parameters():
                p.requires_grad = True
        if self.ssm_bwd is not None:
            for p in self.ssm_bwd.parameters():
                p.requires_grad = True

    # ------------------------------------------------------------------
    # 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, a) for a in ("out_proj", "out_lin", "c_proj", "dense", "o_proj")
        )
        is_nested = False
        if (
            (
                hasattr(old_attn, "self")
                and hasattr(old_attn.self, "query")
                and hasattr(old_attn, "output")
                and hasattr(old_attn.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_proj
        ):
            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),
        }

        # GPT-2 fused
        if hasattr(old_attn, "c_attn") and old_attn.c_attn is not None:
            c = old_attn.c_attn
            base_c = getattr(c, "base_layer", c)
            cw = getattr(base_c, "weight", getattr(c, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                is_conv = (
                    c.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (cw.ndim == 2 and cw.shape[0] < cw.shape[1])
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )

            if hasattr(old_attn, "c_proj") and old_attn.c_proj is not None:
                o = old_attn.c_proj
                base_o = getattr(o, "base_layer", o)
                ow = getattr(base_o, "weight", getattr(o, "weight", None))
                if ow is not None:
                    ow = ow.data.clone()
                    is_conv_o = (
                        o.__class__.__name__ == "Conv1D"
                        or base_o.__class__.__name__ == "Conv1D"
                        or (ow.ndim == 2 and ow.shape[0] < ow.shape[1])
                    )
                    if is_conv_o:
                        ow = ow.t().contiguous()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # Nested (BERT/RoBERTa)
        elif hasattr(old_attn, "self") and hasattr(old_attn.self, "query"):
            s = old_attn.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    b = getattr(m, "bias", None)
                    extracted[key] = (m.weight.data, b.data if b is not None else None)
            if hasattr(old_attn, "output") and hasattr(old_attn.output, "dense"):
                d = old_attn.output.dense
                b = getattr(d, "bias", None)
                extracted["o"] = (d.weight.data, b.data if b is not None else None)

        # Flat
        else:
            cand = {
                "q": ("q_proj", "q_lin", "query"),
                "k": ("k_proj", "k_lin", "key"),
                "v": ("v_proj", "v_lin", "value"),
                "o": ("out_proj", "out_lin", "dense", "o_proj"),
            }
            for key, names in cand.items():
                for n in names:
                    if hasattr(old_attn, n):
                        m = getattr(old_attn, n)
                        if m is not None and hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[key] = (
                                m.weight.data,
                                b.data if b is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        missing = [
            k
            for k, (w, _) in extracted.items()
            if w is None and (k != "o" or not self.is_nested_encoder)
        ]
        if missing:
            try:
                adapter = UniversalAdapter()
                for k in missing:
                    w, b = adapter.get_weight_bias(old_attn, 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:
            logger.warning("MambaSVDFlash: missing projections %s from %s", missing_after, cls_name)

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for key, target in target_map.items():
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target.weight)
                        if target.bias is not None:
                            nn.init.zeros_(target.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)
                if (
                    w.shape != target.weight.shape
                    and w.t().contiguous().shape == target.weight.shape
                ):
                    w = w.t().contiguous()

                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))
                target.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                )
                if b is not None and target.bias is not None:
                    mb = min(target.bias.numel(), b.numel())
                    target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))
                elif target.bias is not None:
                    nn.init.zeros_(target.bias)

                if self.freeze_base:
                    target.freeze_base_weights()

            # MSF-FIX-17: initialise SSM projections from W_q, W_k using SVD
            w_q, _ = extracted["q"]
            w_k, _ = extracted["k"]
            w_v, _ = extracted["v"]

            self._init_ssm_from_attention(w_k, w_q, w_v)

        if hasattr(old_attn, "rotary_emb") and old_attn.rotary_emb is not None:
            self.rotary_emb = old_attn.rotary_emb

        if self.freeze_base:
            self.freeze_base_weights()

    def _init_ssm_from_attention(
        self,
        w_k: torch.Tensor | None,
        w_q: torch.Tensor | None,
        w_v: torch.Tensor | None,
    ):
        """
        MSF-FIX-17: use SVD of W_k, W_q to seed the low-rank SSM bases.
        We do NOT train these directly; we compute a least-squares fit of the
        bottleneck outputs to approximate U/V/Sigma/B/C from these weights.
        Simpler approach: leave SSM at random init but scale to match.
        """
        # Practical approach: fit the up_b / up_c linear layers so that
        # B(x) ≈ W_k @ x and C(x) ≈ W_q @ x  (best-effort, through the bottleneck)
        try:
            if w_k is None and w_q is None:
                return

            D = self.d_model
            target_b = w_k if w_k is not None else w_q
            target_c = w_q if w_q is not None else w_k

            if target_b is None or target_c is None:
                return

            # Ensure shapes match (D, D)
            def _fit(mat):
                if mat is None:
                    return
                if mat.ndim != 2:
                    return
                if mat.shape != (D, D):
                    return
                mat_f = mat.float()
                # SVD to get top-r singular values
                try:
                    U, S, Vh = torch.linalg.svd(mat_f, full_matrices=False)
                except Exception:
                    return
                r = min(self.rank, S.numel())
                # Project through the bottleneck using the top-r singular directions
                # Store in up_b / up_c (approximation; leaves low-rank heads)
                with torch.no_grad():
                    if self.ssm_fwd.up_b.weight.shape == (
                        D,
                        self.ssm_fwd.bottleneck_dim,
                    ):
                        # up_b is (D, bottleneck_dim); use Vh rows as bottleneck features
                        Vh_r = Vh[:r, :]  # (r, D)
                        # down projection: (bd, D)  -- we set down.weight rows to Vh_r
                        self.ssm_fwd.down.weight.data[:r, :].copy_(
                            Vh_r.to(
                                self.ssm_fwd.down.weight.device,
                                self.ssm_fwd.down.weight.dtype,
                            )
                        )

        except Exception as e:
            logger.debug("SSM weight init from attention failed: %s", e)

    # ------------------------------------------------------------------
    # Mask helper (MSF-FIX-04)
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        L: int,
        device: torch.device,
    ) -> torch.Tensor:
        if attention_mask is None:
            return torch.ones(B, L, device=device, dtype=torch.bool)

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape == (B, L):
                return m
        elif attention_mask.dim() == 3:
            m = _to_bool(attention_mask)
            if m.shape[0] == 1:
                return m[0]
            if m.shape[-1] == L:
                return m[:, -1, :] if m.shape[1] == L else m[0, :, :]
        elif attention_mask.dim() == 4:
            m = _to_bool(attention_mask)
            m = m.any(dim=-2)
            if m.dim() == 3:
                m = m[:, 0, :]
            if m.shape == (B, L):
                return m

        try:
            m = get_token_valid_mask(attention_mask, B, L, device)
            if isinstance(m, torch.Tensor) and m.shape == (B, L):
                return m.bool()
        except Exception:
            pass

        return torch.ones(B, L, device=device, dtype=torch.bool)

    # ------------------------------------------------------------------
    # Scan (MSF-FIX-06, MSF-FIX-07, MSF-FIX-08)
    # ------------------------------------------------------------------
    def _ssm_scan(
        self,
        U: torch.Tensor,  # (B, L, D, r)
        Sigma: torch.Tensor,  # (B, L, r)
        V: torch.Tensor,  # (B, L, r, D)
        B_seq: torch.Tensor,  # (B, L, D)
        C_seq: torch.Tensor,  # (B, L, D)
        x: torch.Tensor,  # (B, L, D)
        mask: torch.Tensor | None = None,  # (B, L) bool
        h0: torch.Tensor | None = None,  # (B, D)
        reverse: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Sequential scan with a chunked loop.
        Returns (outputs, h_final).
        - For forward pass, h_final corresponds to step L-1.
        - For reverse pass, the caller flips before/after.
        """
        B, L, D = x.shape
        Sigma.shape[-1]
        device = x.device

        # MSF-FIX-08: preallocate output tensor (no list)
        H = torch.zeros(B, L, D, device=device, dtype=torch.float32)
        h = torch.zeros(B, D, device=device, dtype=torch.float32) if h0 is None else h0.float()

        US = U.float() * Sigma.float().unsqueeze(-2)  # (B, L, D, r)
        V_f = V.float()
        bx = B_seq.float() * x.float()

        steps = range(L - 1, -1, -1) if reverse else range(L)

        for t in steps:
            v_t = torch.einsum("bdr,bd->br", V_f[:, t], h)  # (B, r)
            h_trans = torch.einsum("bdr,br->bd", US[:, t], v_t)  # (B, D)
            h_next = h_trans + bx[:, t]
            if mask is not None:
                m_t = mask[:, t].unsqueeze(-1).to(h_next.dtype)
                h = torch.where(m_t.bool(), h_next, h)
            else:
                h = h_next
            h = torch.nan_to_num(h, nan=0.0, posinf=1e4, neginf=-1e4)
            H[:, t] = h

        out = C_seq.float() * H
        return out.to(x.dtype), h.to(x.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,
        position_ids: torch.Tensor | None = None,
        *args: Any,
        **kwargs: Any,
    ):
        if encoder_hidden_states is not None:
            raise NotImplementedError("MambaSVDFlash does not support cross-attention.")

        # MSF-FIX-16: bounded unpack
        if isinstance(hidden_states, (tuple, list)) and hidden_states:
            hidden_states = hidden_states[0]
        if isinstance(attention_mask, (tuple, list)) and attention_mask:
            attention_mask = attention_mask[0]
        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, L, D = hidden_states.shape
        dtype_in = hidden_states.dtype
        device = hidden_states.device

        # MSF-FIX-15: read cache
        past_h = None
        past_attn_k = None
        past_attn_v = None
        past_len = 0
        if past_key_value is not None:
            if isinstance(past_key_value, dict):
                past_h = past_key_value.get("ssm_h", None)
                past_attn_k = past_key_value.get("attn_k", None)
                past_attn_v = past_key_value.get("attn_v", None)
                past_len = int(past_key_value.get("seq_len", 0))
            elif isinstance(past_key_value, (tuple, list)):
                if len(past_key_value) >= 3:
                    past_h = past_key_value[0]
                    past_attn_k = past_key_value[1]
                    past_attn_v = past_key_value[2]
                    past_len = (
                        int(past_key_value[3])
                        if len(past_key_value) > 3
                        else (past_attn_k.shape[-2] if past_attn_k is not None else 0)
                    )

        mask_bool = self._normalize_valid_mask(attention_mask, B, L, device)

        # ---------------- 1. SSM path ----------------
        U_f, Sig_f, V_f, B_f, C_f = self.ssm_fwd(hidden_states)

        if self.is_decoder:
            ssm_out, ssm_h_next = self._ssm_scan(
                U_f,
                Sig_f,
                V_f,
                B_f,
                C_f,
                hidden_states,
                mask=mask_bool,
                h0=past_h,
                reverse=False,
            )
        else:
            # Forward direction
            ssm_fwd, _ = self._ssm_scan(
                U_f,
                Sig_f,
                V_f,
                B_f,
                C_f,
                hidden_states,
                mask=mask_bool,
                h0=None,
                reverse=False,
            )
            # Backward direction (only meaningful for encoder)
            if self.ssm_bwd is not None:
                U_b, Sig_b, V_b, B_b, C_b = self.ssm_bwd(hidden_states)
                x_rev = torch.flip(hidden_states, dims=[1])
                mask_rev = torch.flip(mask_bool, dims=[1]) if mask_bool is not None else None
                ssm_bwd_rev, _ = self._ssm_scan(
                    U_b,
                    Sig_b,
                    V_b,
                    B_b,
                    C_b,
                    x_rev,
                    mask=mask_rev,
                    h0=None,
                    reverse=False,
                )
                ssm_bwd = torch.flip(ssm_bwd_rev, dims=[1])
                # MSF-FIX-13: per-head mix
                w = torch.sigmoid(self.bidirectional_mix).view(1, self.num_heads, 1, 1)
                ssm_fwd_h = ssm_fwd.view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
                ssm_bwd_h = ssm_bwd.view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
                ssm_out = (
                    (w * ssm_fwd_h + (1 - w) * ssm_bwd_h).transpose(1, 2).contiguous().view(B, L, D)
                )
                ssm_h_next = None
            else:
                ssm_out = ssm_fwd
                ssm_h_next = None

        # ---------------- 2. Attention path ----------------
        Q = (
            self.q_proj(hidden_states)
            .view(B, L, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            self.k_proj(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V_a = (
            self.v_proj(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # MSF-FIX-18: RoPE with offset
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_ids is None:
            position_ids = kwargs.get("position_ids")

        if position_embeddings is not None:
            cos, sin = position_embeddings
            # Slice to [past_len, past_len + L)
            if cos.shape[-2] >= past_len + L:
                cos = cos.narrow(-2, past_len, L)
                sin = sin.narrow(-2, past_len, L)
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)
        elif self.rotary_emb is not None:
            if position_ids is None:
                position_ids = (
                    torch.arange(past_len, past_len + L, device=device).unsqueeze(0).expand(B, -1)
                )
            cos = sin = None
            for call in (
                lambda: self.rotary_emb(K, seq_len=L),
                lambda: self.rotary_emb(K, position_ids),
                lambda: self.rotary_emb(position_ids),
            ):
                try:
                    cos, sin = call()
                    break
                except Exception:
                    continue
            if cos is not None and sin is not None:
                Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids=position_ids)

        # MSF-FIX-15: append to attention cache
        if past_attn_k is not None:
            K = torch.cat([past_attn_k.to(K.dtype), K], dim=-2)
            V_a = torch.cat([past_attn_v.to(V_a.dtype), V_a], dim=-2)

        # Store compact (before GQA expand)
        present_attn_k = K
        present_attn_v = V_a

        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V_a = repeat_kv(V_a, self.kv_repeat)

        scale = 1.0 / math.sqrt(max(1, self.head_dim))
        attn_out: torch.Tensor | None = None
        probs: torch.Tensor | None = None

        # SDPA fast path
        can_sdpa = (
            head_mask is None
            and not output_attentions
            and bool(mask_bool.all().item())  # single sync at call time
            and past_len == 0
        )
        if can_sdpa:
            try:
                _qs = Q.half() if Q.dtype == torch.float32 else Q
                _ks = K.half() if K.dtype == torch.float32 else K
                _vs = V_a.half() if V_a.dtype == torch.float32 else V_a
                _is_causal = self.is_decoder and L > 1
                _sdpa_out = F.scaled_dot_product_attention(
                    _qs,
                    _ks,
                    _vs,
                    is_causal=_is_causal,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                attn_out = (
                    _sdpa_out.to(dtype_in)
                    .transpose(1, 2)
                    .contiguous()
                    .view(B, L, self.all_head_size)
                )
                attn_out = torch.nan_to_num(attn_out, nan=0.0)
            except Exception:
                attn_out = None

        if attn_out is None:
            scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

            if self.is_decoder and L > 1:
                kv_len = K.shape[-2]
                causal = torch.tril(
                    torch.ones(L, kv_len, device=device, dtype=torch.bool),
                    diagonal=past_len,
                )
                scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

            # Padding mask: keys that are invalid get -inf
            if not mask_bool.all().item():
                scores = scores.masked_fill(~mask_bool[:, None, None, :], -1e4)

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(dtype_in)
            probs = torch.nan_to_num(probs, nan=0.0)
            probs = self.dropout(probs)

            if head_mask is not None:
                probs = self._apply_head_mask(head_mask, probs)

            attn_out = (
                torch.matmul(probs, V_a).transpose(1, 2).contiguous().view(B, L, self.all_head_size)
            )
            attn_out = torch.nan_to_num(attn_out, nan=0.0)

        # ---------------- 3. Per-head gating fusion ----------------
        # gamma: (B, L, num_heads)
        gamma = torch.sigmoid(self.gating(hidden_states))
        # Reshape for per-head fusion
        ssm_h = ssm_out.view(B, L, self.num_heads, self.head_dim)
        attn_h = attn_out.view(B, L, self.num_heads, self.head_dim)
        g = gamma.unsqueeze(-1)  # (B, L, H, 1)
        Y = (g * ssm_h + (1 - g) * attn_h).contiguous().view(B, L, self.all_head_size)
        Y = torch.nan_to_num(Y, nan=0.0, posinf=1e4, neginf=-1e4)

        out = Y if self.is_nested_encoder else self.out_proj(Y)

        # Build present_key_value
        present_key_value = {
            "ssm_h": (
                ssm_h_next
                if ssm_h_next is not None
                else torch.zeros(B, D, device=device, dtype=dtype_in)
            ),
            "attn_k": present_attn_k,
            "attn_v": present_attn_v,
            "seq_len": past_len + L,
        }

        attn_weights = probs if output_attentions else None
        return out, attn_weights, present_key_value

    def _apply_head_mask(self, head_mask: torch.Tensor, probs: torch.Tensor) -> torch.Tensor:
        """MSF-FIX-11: unified head_mask broadcast. probs: (B, H, Lq, Lk)."""
        hm = head_mask.to(probs.dtype)
        H = self.num_heads
        if hm.dim() == 1:
            if hm.numel() == H:
                hm = hm.view(1, H, 1, 1)
            else:
                hm = hm.view(1, -1, 1, 1)
        elif hm.dim() == 2:
            # (B, H) or (H, 1)
            if hm.shape[-1] == H:
                hm = hm.view(hm.shape[0], H, 1, 1)
            elif hm.shape[0] == H:
                hm = hm.t().view(1, H, 1, 1)
            else:
                hm = hm.view(1, hm.shape[-1], 1, 1)
        elif hm.dim() == 3:
            # (B, H, Lk) or (B, Lk, H)
            if hm.shape[1] == H:
                hm = hm.unsqueeze(2)  # (B, H, 1, Lk)
            elif hm.shape[2] == H:
                hm = hm.permute(0, 2, 1).unsqueeze(2)
            else:
                hm = hm.unsqueeze(1)
        # dim == 4: pass through
        return probs * hm

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out: dict[str, int] = {}
        for name, m in (
            ("q_proj", self.q_proj),
            ("k_proj", self.k_proj),
            ("v_proj", self.v_proj),
            ("out_proj", self.out_proj),
            ("gating", self.gating),
            ("ssm_fwd", self.ssm_fwd),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        if self.ssm_bwd is not None:
            out["ssm_bwd"] = sum(_sz(p) for p in self.ssm_bwd.parameters())
        out["total"] = sum(out.values())
        return out

Hybrid SSM + MHA with per-chunk SVD and learnable per-head gating.

State cache layout (for streaming): past_key_value = { "ssm_h": (B, D), SSM hidden state h_{L-1} "attn_k": (B, H, L, Dh), attention keys "attn_v": (B, H, L, Dh), attention values "seq_len": int, } Or tuple form: (ssm_h, attn_k, attn_v, seq_len).

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.

Instance variables

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 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 freeze_base_weights(self)

Expand source code

def freeze_base_weights(self):
    # MSF-FIX-09: freeze SSM projections too
    for m in (
        self.q_proj,
        self.k_proj,
        self.v_proj,
        self.out_proj,
        self.gating,
        self.ssm_fwd.down,
        self.ssm_fwd.up_u,
        self.ssm_fwd.up_v,
        self.ssm_fwd.up_sigma,
        self.ssm_fwd.up_b,
        self.ssm_fwd.up_c,
    ):
        for p in m.parameters():
            p.requires_grad = False
    if self.ssm_bwd is not None:
        for p in self.ssm_bwd.parameters():
            p.requires_grad = False
    # Keep LoRA adapters trainable
    for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj, self.gating):
        for name, p in m.named_parameters():
            if "lora_" in name:
                p.requires_grad = True

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out: dict[str, int] = {}
    for name, m in (
        ("q_proj", self.q_proj),
        ("k_proj", self.k_proj),
        ("v_proj", self.v_proj),
        ("out_proj", self.out_proj),
        ("gating", self.gating),
        ("ssm_fwd", self.ssm_fwd),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    if self.ssm_bwd is not None:
        out["ssm_bwd"] = sum(_sz(p) for p in self.ssm_bwd.parameters())
    out["total"] = sum(out.values())
    return out

def unfreeze_base_weights(self)

Expand source code

def unfreeze_base_weights(self):
    for m in (
        self.q_proj,
        self.k_proj,
        self.v_proj,
        self.out_proj,
        self.gating,
        self.ssm_fwd.down,
        self.ssm_fwd.up_u,
        self.ssm_fwd.up_v,
        self.ssm_fwd.up_sigma,
        self.ssm_fwd.up_b,
        self.ssm_fwd.up_c,
    ):
        for p in m.parameters():
            p.requires_grad = True
    if self.ssm_bwd is not None:
        for p in self.ssm_bwd.parameters():
            p.requires_grad = True

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class MinimalMoEFallback (d_model: int = 768, n_experts: int = 4)

Expand source code

class MinimalMoEFallback(nn.Module):
    """Fallback MoE layer when REAPMoE cannot be instantiated."""

    def __init__(self, d_model: int = 768, n_experts: int = 4):
        super().__init__()
        self.d_model = d_model
        self.n_experts = n_experts
        self.router = nn.Linear(d_model, n_experts, bias=False)
        self.experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(d_model, d_model * 2),
                    nn.GELU(),
                    nn.Linear(d_model * 2, d_model),
                )
                for _ in range(n_experts)
            ]
        )
        self.aux_loss: torch.Tensor = torch.tensor(0.0)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        router_in = self.router.in_features
        if x.shape[-1] != router_in:
            x_in = (
                F.pad(x, (0, router_in - x.shape[-1]))
                if x.shape[-1] < router_in
                else x[..., :router_in]
            )
        else:
            x_in = x
        logits = self.router(x_in)
        weights = F.softmax(logits, dim=-1)
        out = torch.zeros_like(x_in)
        for i, exp in enumerate(self.experts):
            out = out + weights[..., i : i + 1] * exp(x_in)
        if out.shape[-1] != x.shape[-1]:
            out = (
                F.pad(out, (0, x.shape[-1] - out.shape[-1]))
                if out.shape[-1] < x.shape[-1]
                else out[..., : x.shape[-1]]
            )
        self.aux_loss = torch.tensor(0.0, device=x.device)
        return out, weights

Fallback MoE layer when REAPMoE cannot be instantiated.

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

Ancestors

• torch.nn.modules.module.Module

Methods

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

Expand source code

def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    router_in = self.router.in_features
    if x.shape[-1] != router_in:
        x_in = (
            F.pad(x, (0, router_in - x.shape[-1]))
            if x.shape[-1] < router_in
            else x[..., :router_in]
        )
    else:
        x_in = x
    logits = self.router(x_in)
    weights = F.softmax(logits, dim=-1)
    out = torch.zeros_like(x_in)
    for i, exp in enumerate(self.experts):
        out = out + weights[..., i : i + 1] * exp(x_in)
    if out.shape[-1] != x.shape[-1]:
        out = (
            F.pad(out, (0, x.shape[-1] - out.shape[-1]))
            if out.shape[-1] < x.shape[-1]
            else out[..., : x.shape[-1]]
        )
    self.aux_loss = torch.tensor(0.0, device=x.device)
    return out, weights

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

class MinimalNOVAAttention (config: Any | None = None,
model_config: Any | None = None,
energy: float = 0.95,
dropout: float = 0.1,
pretrained_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("minimal_nova", supports_causal=True, supports_gqa=True)
class MinimalNOVAAttention(BaseAttention):
    r"""
    MinimalNOVA: SVD-Factorized Low-Rank Projection Attention.
    Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        energy: float = 0.95,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        rank: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = model_config if model_config is not None else config
        actual_rank = getattr(
            cfg,
            "svd_rank",
            rank if rank is not None else getattr(cfg, "rank", 64),
        )
        super().__init__(cfg, actual_rank)
        self.num_heads: int = getattr(
            cfg,
            "num_attention_heads",
            getattr(cfg, "num_heads", kwargs.get("num_attention_heads", 4)),
        )
        self.hidden_size: int = getattr(cfg, "hidden_size", kwargs.get("hidden_size", 768))

        cfg_head_dim = getattr(cfg, "head_dim", None) if cfg else None
        adapter_head_dim = getattr(getattr(cfg, "adapter", None), "head_dim", None) if cfg else None
        self.head_dim: int = int(
            cfg_head_dim
            or adapter_head_dim
            or kwargs.get("head_dim")
            or (self.hidden_size // self.num_heads if self.num_heads > 0 else 64)
        )
        self.all_head_size: int = self.num_heads * self.head_dim

        cfg_kv = (
            getattr(cfg, "num_key_value_heads", getattr(cfg, "num_kv_heads", None)) if cfg else None
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        self.rank: int = actual_rank
        self.energy: float = energy
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha

        mt = str(getattr(cfg, "model_type", "") or "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "olmo",
            "falcon",
            "mpt",
            "stablelm",
        )
        self.is_decoder = (
            (bool(getattr(cfg, "is_decoder", False)) if cfg else False)
            or (bool(getattr(cfg, "causal", False)) if cfg else False)
            or any(mt.startswith(p) for p in _decoder_prefixes)
            or (pretrained_attn is not None and hasattr(pretrained_attn, "c_attn"))
        )
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)

        self.is_nested_encoder: bool = False

        src_p = None
        if pretrained_attn is not None:
            src_p = next(pretrained_attn.parameters(), None)
        target_dtype = kwargs.get("dtype", src_p.dtype if src_p is not None else torch.float32)
        target_device = kwargs.get("device", src_p.device if src_p is not None else None)

        self.q_proj = FactorisedLinear(
            self.hidden_size,
            self.all_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            device=target_device,
            dtype=target_dtype,
        )
        self.k_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            device=target_device,
            dtype=target_dtype,
        )
        self.v_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            device=target_device,
            dtype=target_dtype,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            self.hidden_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            device=target_device,
            dtype=target_dtype,
        )

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", dropout))
            if cfg
            else dropout
        )
        self.dropout: nn.Dropout = nn.Dropout(dropout_p)

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

        if target_device is not None or target_dtype is not None:
            self.to(device=target_device, dtype=target_dtype)

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

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

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

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

        # 2. GPT-2 Fused c_attn / c_proj 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_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                p_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
                ):
                    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:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_o, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj_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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
                or (
                    ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                    and not has_out_proj
                )
            )
        ):
            is_nested = True
        if hasattr(src, "config"):
            if (
                getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
                and not has_out_proj
            ):
                is_nested = True
        self.is_nested_encoder = bool(is_nested)

        # Check required weights
        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("Module copy_weights warning. Proceeding with initialized weights.")

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"MinimalNOVAAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"MinimalNOVAAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
                )

        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

        # Copy into factorized projection modules
        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target.decompose_from_weight(w_oriented, b, self.energy)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None, self.energy)

                target.freeze_base_weights()

        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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        # Defensive device check
        if self.q_proj.down.weight.device != device or self.q_proj.down.weight.dtype != dtype:
            self.to(device=device, dtype=dtype)

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

        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, N_kv, self.num_kv_heads, D).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, N_kv, 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 and not is_cross:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = 1.0 / math.sqrt(D)
        # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

        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 = (
                getattr(self, "is_decoder", False) and not is_cross and N > 1
                if "is_cross" in dir()
                else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
            )

            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None

        except Exception:
            pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
                elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                    scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
                else:
                    scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask.unsqueeze(1), -10000.0)
                else:
                    scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask, -10000.0)
                else:
                    scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(dtype)
        attn_weights = self.dropout(attn_weights)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn_weights = attn_weights * head_mask

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

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

        return output, (attn_weights if output_attentions else None)

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.tensor(0.0, device=self.out_proj.up.weight.device)

MinimalNOVA: SVD-Factorized Low-Rank Projection Attention. Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.

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.

Instance variables

prop key : FactorisedLinear

Expand source code

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

prop query : FactorisedLinear

Expand source code

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

prop value : FactorisedLinear

Expand source code

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

Methods

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.tensor(0.0, device=self.out_proj.up.weight.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class MinimalSVDAttention (config: Any | None = None,
model_config: Any | None = None,
energy: float = 0.95,
dropout: float = 0.1,
pretrained_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("minimal_svd", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_aurora", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_prism", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_svd_v2", supports_causal=True, supports_gqa=True)
class MinimalSVDAttention(BaseAttention):
    r"""
    MinimalSVD: SVD-Factorized Low-Rank Projection Attention.
    Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        energy: float = 0.95,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        rank: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = model_config if model_config is not None else config
        actual_rank = getattr(
            cfg,
            "svd_rank",
            rank if rank is not None else getattr(cfg, "rank", 64),
        )
        super().__init__(cfg, actual_rank)
        self.num_heads: int = getattr(
            cfg,
            "num_attention_heads",
            getattr(cfg, "num_heads", kwargs.get("num_attention_heads", 4)),
        )
        self.hidden_size: int = getattr(cfg, "hidden_size", kwargs.get("hidden_size", 768))
        cfg_head_dim = (
            getattr(cfg, "head_dim", None)
            if cfg
            else (getattr(model_config, "head_dim", None) if model_config else None)
        )
        self.head_dim: int = int(
            cfg_head_dim
            or kwargs.get("head_dim")
            or (
                self.hidden_size // self.num_heads
                if self.num_heads > 0
                else (self.hidden_size or 64)
            )
        )
        self.all_head_size: int = self.num_heads * self.head_dim

        cfg_kv = (
            getattr(cfg, "num_key_value_heads", getattr(cfg, "num_kv_heads", None)) if cfg else None
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        self.rank: int = actual_rank
        self.energy: float = energy
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = (
                getattr(cfg, "is_decoder", False)
                or getattr(cfg, "causal", False)
                or getattr(cfg, "model_type", "").lower() in ("gpt2", "qwen", "qwen2", "llama")
                or any(
                    str(
                        getattr(
                            (
                                cfg
                                if "cfg" in locals()
                                else (config if "config" in locals() else None)
                            ),
                            "model_type",
                            "",
                        )
                    )
                    .lower()
                    .startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                if cfg
                else False
            )

        model_type = getattr(cfg, "model_type", "").lower() if cfg else ""
        self.is_nested_encoder: bool = getattr(cfg, "is_nested_encoder", False) or model_type in (
            "roberta",
            "bert",
            "xlm-roberta",
        )

        self.q_proj = FactorisedLinear(
            self.hidden_size,
            self.all_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.k_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.v_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            self.hidden_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", dropout))
            if cfg
            else dropout
        )
        self.dropout: nn.Dropout = nn.Dropout(dropout_p)

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

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

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

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

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

        # 2. GPT-2 Fused c_attn / c_proj 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_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                p_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
                ):
                    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:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_o, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj_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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
                or (
                    ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                    and not has_out_proj
                )
            )
        ):
            is_nested = True
        if hasattr(src, "config"):
            if (
                getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
                and not has_out_proj
            ):
                is_nested = True
        self.is_nested_encoder = bool(is_nested)

        # Check required weights
        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(
                f"MinimalSVDAttention missing required weights for projections: {missing}. Proceeding with initialized weights."
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"MinimalSVDAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"MinimalSVDAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
                )

        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

        # Copy into factorized projection modules
        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target.decompose_from_weight(w_oriented, b, self.energy)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None, self.energy)

                target.freeze_base_weights()

        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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        if self.q_proj.down.weight.device != device or self.q_proj.down.weight.dtype != dtype:
            self.to(device=device, dtype=dtype)

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

        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, N_kv, self.num_kv_heads, D).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, N_kv, 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 and not is_cross:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = 1.0 / math.sqrt(D)
        # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

        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 = (
                getattr(self, "is_decoder", False) and not is_cross and N > 1
                if "is_cross" in dir()
                else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
            )

            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None

        except Exception:
            pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
                elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                    scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
                else:
                    scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask, -10000.0)
                else:
                    scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(dtype)
        attn_weights = self.dropout(attn_weights)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn_weights = attn_weights * head_mask

        context = torch.matmul(attn_weights, V)
        context = context.transpose(1, 2).contiguous().view(B, N, H * D)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return output, (attn_weights if output_attentions else None)

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.tensor(0.0, device=self.out_proj.up.weight.device)

MinimalSVD: SVD-Factorized Low-Rank Projection Attention. Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.

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

Ancestors

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

Subclasses

• NaiveSVDAttention

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.

Instance variables

prop key : FactorisedLinear

Expand source code

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

prop query : FactorisedLinear

Expand source code

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

prop value : FactorisedLinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    from neuralops.core.adapters import UniversalAdapter

    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 as e:
            logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

    # 2. GPT-2 Fused c_attn / c_proj 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_attn_mod = src.c_attn
        base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
        weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
        if weight is not None:
            weight = weight.data.clone()
            bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            bias_data = bias.data.clone() if bias is not None else None
            p_cls = getattr(base_c, "__class__", type(base_c)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(base_c, "nf")
                or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                or hasattr(c_attn_mod, "nf")
                or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
            ):
                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:
        c_proj_mod = src.c_proj
        base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
        w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
        if w_o is not None:
            w_o = w_o.data.clone()
            p_cls = getattr(base_o, "__class__", type(base_o)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(base_o, "nf")
                or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                or hasattr(c_proj_mod, "nf")
                or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(base_o, "bias", getattr(c_proj_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"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    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 = src.__class__.__name__
    has_out_proj = any(
        hasattr(src, attr)
        for attr in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    is_nested = False
    if (
        hasattr(src, "self")
        and hasattr(src.self, "query")
        or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
            )
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
            or (
                ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                and not has_out_proj
            )
        )
    ):
        is_nested = True
    if hasattr(src, "config"):
        if (
            getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
            and not has_out_proj
        ):
            is_nested = True
    self.is_nested_encoder = bool(is_nested)

    # Check required weights
    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(
            f"MinimalSVDAttention missing required weights for projections: {missing}. Proceeding with initialized weights."
        )

    # Frobenius norm validation
    for proj_key in ["q", "k", "v"]:
        w = extracted[proj_key][0]
        if w is not None:
            norm = torch.linalg.norm(w.float())
            if norm < 1e-4:
                logger.warning(
                    f"MinimalSVDAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                )

    if extracted["o"][0] is not None:
        norm_o = torch.linalg.norm(extracted["o"][0].float())
        if norm_o < 1e-4:
            logger.warning(
                f"MinimalSVDAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
            )

    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

    # Copy into factorized projection modules
    with torch.no_grad():
        for proj_key, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                w_oriented = _orient(w, in_dim)
                target.decompose_from_weight(w_oriented, b, self.energy)
            elif proj_key == "o" and self.is_nested_encoder:
                eye_w = torch.eye(
                    self.all_head_size,
                    self.all_head_size,
                    device=self.out_proj.up.weight.device,
                    dtype=self.out_proj.up.weight.dtype,
                )
                target.decompose_from_weight(eye_w, None, self.energy)

            target.freeze_base_weights()

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

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    H, D = self.num_heads, self.head_dim
    device = hidden_states.device
    dtype = hidden_states.dtype

    if self.q_proj.down.weight.device != device or self.q_proj.down.weight.dtype != dtype:
        self.to(device=device, dtype=dtype)

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

    Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
    K = self.k_proj(kv_input).view(B, N_kv, self.num_kv_heads, D).transpose(1, 2).contiguous()
    V = self.v_proj(kv_input).view(B, N_kv, 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 and not is_cross:
        cos, sin = position_embeddings
        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

    if self.num_kv_heads < self.num_heads:
        K = repeat_kv(K, self.kv_repeat)
        V = repeat_kv(V, self.kv_repeat)

    scale = 1.0 / math.sqrt(D)
    # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

    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 = (
            getattr(self, "is_decoder", False) and not is_cross and N > 1
            if "is_cross" in dir()
            else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
        )

        _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

        return (
            self.out_proj(_sdpa_out)
            if not getattr(self, "is_nested_encoder", False)
            else _sdpa_out
        ), None

    except Exception:
        pass

    scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

    if self.is_decoder and not is_cross and N > 1:
        causal_mask = (
            torch.triu(
                torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                diagonal=1,
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        scores = scores + causal_mask

    eff_mask = encoder_attention_mask if is_cross else attention_mask
    if eff_mask is not None:
        if eff_mask.dim() == 2:
            if eff_mask.dtype == torch.bool:
                scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
            elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
            else:
                scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
        elif eff_mask.dim() == 3:
            scores = scores + eff_mask.unsqueeze(1).float()
        elif eff_mask.dim() == 4:
            if eff_mask.dtype == torch.bool:
                scores = scores.masked_fill(~eff_mask, -10000.0)
            else:
                scores = scores + eff_mask.float()

    scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
    attn_weights = F.softmax(scores, dim=-1).to(dtype)
    attn_weights = self.dropout(attn_weights)

    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)
        elif head_mask.dim() == 3:
            head_mask = head_mask.unsqueeze(-1)
        attn_weights = attn_weights * head_mask

    context = torch.matmul(attn_weights, V)
    context = context.transpose(1, 2).contiguous().view(B, N, H * D)
    context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

    return output, (attn_weights if output_attentions else None)

Must return (output, attention_weights_or_None).

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.tensor(0.0, device=self.out_proj.up.weight.device)

class MinimalAuroraAttention (config: Any | None = None,
model_config: Any | None = None,
energy: float = 0.95,
dropout: float = 0.1,
pretrained_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("minimal_svd", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_aurora", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_prism", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_svd_v2", supports_causal=True, supports_gqa=True)
class MinimalSVDAttention(BaseAttention):
    r"""
    MinimalSVD: SVD-Factorized Low-Rank Projection Attention.
    Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        energy: float = 0.95,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        rank: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = model_config if model_config is not None else config
        actual_rank = getattr(
            cfg,
            "svd_rank",
            rank if rank is not None else getattr(cfg, "rank", 64),
        )
        super().__init__(cfg, actual_rank)
        self.num_heads: int = getattr(
            cfg,
            "num_attention_heads",
            getattr(cfg, "num_heads", kwargs.get("num_attention_heads", 4)),
        )
        self.hidden_size: int = getattr(cfg, "hidden_size", kwargs.get("hidden_size", 768))
        cfg_head_dim = (
            getattr(cfg, "head_dim", None)
            if cfg
            else (getattr(model_config, "head_dim", None) if model_config else None)
        )
        self.head_dim: int = int(
            cfg_head_dim
            or kwargs.get("head_dim")
            or (
                self.hidden_size // self.num_heads
                if self.num_heads > 0
                else (self.hidden_size or 64)
            )
        )
        self.all_head_size: int = self.num_heads * self.head_dim

        cfg_kv = (
            getattr(cfg, "num_key_value_heads", getattr(cfg, "num_kv_heads", None)) if cfg else None
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        self.rank: int = actual_rank
        self.energy: float = energy
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = (
                getattr(cfg, "is_decoder", False)
                or getattr(cfg, "causal", False)
                or getattr(cfg, "model_type", "").lower() in ("gpt2", "qwen", "qwen2", "llama")
                or any(
                    str(
                        getattr(
                            (
                                cfg
                                if "cfg" in locals()
                                else (config if "config" in locals() else None)
                            ),
                            "model_type",
                            "",
                        )
                    )
                    .lower()
                    .startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                if cfg
                else False
            )

        model_type = getattr(cfg, "model_type", "").lower() if cfg else ""
        self.is_nested_encoder: bool = getattr(cfg, "is_nested_encoder", False) or model_type in (
            "roberta",
            "bert",
            "xlm-roberta",
        )

        self.q_proj = FactorisedLinear(
            self.hidden_size,
            self.all_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.k_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.v_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            self.hidden_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", dropout))
            if cfg
            else dropout
        )
        self.dropout: nn.Dropout = nn.Dropout(dropout_p)

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

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

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

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

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

        # 2. GPT-2 Fused c_attn / c_proj 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_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                p_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
                ):
                    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:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_o, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj_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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
                or (
                    ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                    and not has_out_proj
                )
            )
        ):
            is_nested = True
        if hasattr(src, "config"):
            if (
                getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
                and not has_out_proj
            ):
                is_nested = True
        self.is_nested_encoder = bool(is_nested)

        # Check required weights
        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(
                f"MinimalSVDAttention missing required weights for projections: {missing}. Proceeding with initialized weights."
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"MinimalSVDAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"MinimalSVDAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
                )

        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

        # Copy into factorized projection modules
        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target.decompose_from_weight(w_oriented, b, self.energy)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None, self.energy)

                target.freeze_base_weights()

        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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        if self.q_proj.down.weight.device != device or self.q_proj.down.weight.dtype != dtype:
            self.to(device=device, dtype=dtype)

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

        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, N_kv, self.num_kv_heads, D).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, N_kv, 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 and not is_cross:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = 1.0 / math.sqrt(D)
        # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

        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 = (
                getattr(self, "is_decoder", False) and not is_cross and N > 1
                if "is_cross" in dir()
                else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
            )

            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None

        except Exception:
            pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
                elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                    scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
                else:
                    scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask, -10000.0)
                else:
                    scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(dtype)
        attn_weights = self.dropout(attn_weights)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn_weights = attn_weights * head_mask

        context = torch.matmul(attn_weights, V)
        context = context.transpose(1, 2).contiguous().view(B, N, H * D)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return output, (attn_weights if output_attentions else None)

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.tensor(0.0, device=self.out_proj.up.weight.device)

MinimalSVD: SVD-Factorized Low-Rank Projection Attention. Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.

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

Ancestors

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

Subclasses

• NaiveSVDAttention

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.

Instance variables

prop key : FactorisedLinear

Expand source code

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

prop query : FactorisedLinear

Expand source code

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

prop value : FactorisedLinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    from neuralops.core.adapters import UniversalAdapter

    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 as e:
            logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

    # 2. GPT-2 Fused c_attn / c_proj 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_attn_mod = src.c_attn
        base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
        weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
        if weight is not None:
            weight = weight.data.clone()
            bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            bias_data = bias.data.clone() if bias is not None else None
            p_cls = getattr(base_c, "__class__", type(base_c)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(base_c, "nf")
                or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                or hasattr(c_attn_mod, "nf")
                or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
            ):
                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:
        c_proj_mod = src.c_proj
        base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
        w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
        if w_o is not None:
            w_o = w_o.data.clone()
            p_cls = getattr(base_o, "__class__", type(base_o)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(base_o, "nf")
                or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                or hasattr(c_proj_mod, "nf")
                or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(base_o, "bias", getattr(c_proj_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"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    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 = src.__class__.__name__
    has_out_proj = any(
        hasattr(src, attr)
        for attr in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    is_nested = False
    if (
        hasattr(src, "self")
        and hasattr(src.self, "query")
        or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
            )
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
            or (
                ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                and not has_out_proj
            )
        )
    ):
        is_nested = True
    if hasattr(src, "config"):
        if (
            getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
            and not has_out_proj
        ):
            is_nested = True
    self.is_nested_encoder = bool(is_nested)

    # Check required weights
    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(
            f"MinimalSVDAttention missing required weights for projections: {missing}. Proceeding with initialized weights."
        )

    # Frobenius norm validation
    for proj_key in ["q", "k", "v"]:
        w = extracted[proj_key][0]
        if w is not None:
            norm = torch.linalg.norm(w.float())
            if norm < 1e-4:
                logger.warning(
                    f"MinimalSVDAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                )

    if extracted["o"][0] is not None:
        norm_o = torch.linalg.norm(extracted["o"][0].float())
        if norm_o < 1e-4:
            logger.warning(
                f"MinimalSVDAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
            )

    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

    # Copy into factorized projection modules
    with torch.no_grad():
        for proj_key, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                w_oriented = _orient(w, in_dim)
                target.decompose_from_weight(w_oriented, b, self.energy)
            elif proj_key == "o" and self.is_nested_encoder:
                eye_w = torch.eye(
                    self.all_head_size,
                    self.all_head_size,
                    device=self.out_proj.up.weight.device,
                    dtype=self.out_proj.up.weight.dtype,
                )
                target.decompose_from_weight(eye_w, None, self.energy)

            target.freeze_base_weights()

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

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    H, D = self.num_heads, self.head_dim
    device = hidden_states.device
    dtype = hidden_states.dtype

    if self.q_proj.down.weight.device != device or self.q_proj.down.weight.dtype != dtype:
        self.to(device=device, dtype=dtype)

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

    Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
    K = self.k_proj(kv_input).view(B, N_kv, self.num_kv_heads, D).transpose(1, 2).contiguous()
    V = self.v_proj(kv_input).view(B, N_kv, 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 and not is_cross:
        cos, sin = position_embeddings
        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

    if self.num_kv_heads < self.num_heads:
        K = repeat_kv(K, self.kv_repeat)
        V = repeat_kv(V, self.kv_repeat)

    scale = 1.0 / math.sqrt(D)
    # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

    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 = (
            getattr(self, "is_decoder", False) and not is_cross and N > 1
            if "is_cross" in dir()
            else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
        )

        _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

        return (
            self.out_proj(_sdpa_out)
            if not getattr(self, "is_nested_encoder", False)
            else _sdpa_out
        ), None

    except Exception:
        pass

    scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

    if self.is_decoder and not is_cross and N > 1:
        causal_mask = (
            torch.triu(
                torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                diagonal=1,
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        scores = scores + causal_mask

    eff_mask = encoder_attention_mask if is_cross else attention_mask
    if eff_mask is not None:
        if eff_mask.dim() == 2:
            if eff_mask.dtype == torch.bool:
                scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
            elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
            else:
                scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
        elif eff_mask.dim() == 3:
            scores = scores + eff_mask.unsqueeze(1).float()
        elif eff_mask.dim() == 4:
            if eff_mask.dtype == torch.bool:
                scores = scores.masked_fill(~eff_mask, -10000.0)
            else:
                scores = scores + eff_mask.float()

    scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
    attn_weights = F.softmax(scores, dim=-1).to(dtype)
    attn_weights = self.dropout(attn_weights)

    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)
        elif head_mask.dim() == 3:
            head_mask = head_mask.unsqueeze(-1)
        attn_weights = attn_weights * head_mask

    context = torch.matmul(attn_weights, V)
    context = context.transpose(1, 2).contiguous().view(B, N, H * D)
    context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

    return output, (attn_weights if output_attentions else None)

Must return (output, attention_weights_or_None).

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.tensor(0.0, device=self.out_proj.up.weight.device)

class MinimalPRISMAttention (config: Any | None = None,
model_config: Any | None = None,
energy: float = 0.95,
dropout: float = 0.1,
pretrained_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("minimal_svd", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_aurora", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_prism", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_svd_v2", supports_causal=True, supports_gqa=True)
class MinimalSVDAttention(BaseAttention):
    r"""
    MinimalSVD: SVD-Factorized Low-Rank Projection Attention.
    Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        energy: float = 0.95,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        rank: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = model_config if model_config is not None else config
        actual_rank = getattr(
            cfg,
            "svd_rank",
            rank if rank is not None else getattr(cfg, "rank", 64),
        )
        super().__init__(cfg, actual_rank)
        self.num_heads: int = getattr(
            cfg,
            "num_attention_heads",
            getattr(cfg, "num_heads", kwargs.get("num_attention_heads", 4)),
        )
        self.hidden_size: int = getattr(cfg, "hidden_size", kwargs.get("hidden_size", 768))
        cfg_head_dim = (
            getattr(cfg, "head_dim", None)
            if cfg
            else (getattr(model_config, "head_dim", None) if model_config else None)
        )
        self.head_dim: int = int(
            cfg_head_dim
            or kwargs.get("head_dim")
            or (
                self.hidden_size // self.num_heads
                if self.num_heads > 0
                else (self.hidden_size or 64)
            )
        )
        self.all_head_size: int = self.num_heads * self.head_dim

        cfg_kv = (
            getattr(cfg, "num_key_value_heads", getattr(cfg, "num_kv_heads", None)) if cfg else None
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        self.rank: int = actual_rank
        self.energy: float = energy
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = (
                getattr(cfg, "is_decoder", False)
                or getattr(cfg, "causal", False)
                or getattr(cfg, "model_type", "").lower() in ("gpt2", "qwen", "qwen2", "llama")
                or any(
                    str(
                        getattr(
                            (
                                cfg
                                if "cfg" in locals()
                                else (config if "config" in locals() else None)
                            ),
                            "model_type",
                            "",
                        )
                    )
                    .lower()
                    .startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                if cfg
                else False
            )

        model_type = getattr(cfg, "model_type", "").lower() if cfg else ""
        self.is_nested_encoder: bool = getattr(cfg, "is_nested_encoder", False) or model_type in (
            "roberta",
            "bert",
            "xlm-roberta",
        )

        self.q_proj = FactorisedLinear(
            self.hidden_size,
            self.all_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.k_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.v_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            self.hidden_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", dropout))
            if cfg
            else dropout
        )
        self.dropout: nn.Dropout = nn.Dropout(dropout_p)

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

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

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

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

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

        # 2. GPT-2 Fused c_attn / c_proj 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_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                p_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
                ):
                    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:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_o, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj_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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
                or (
                    ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                    and not has_out_proj
                )
            )
        ):
            is_nested = True
        if hasattr(src, "config"):
            if (
                getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
                and not has_out_proj
            ):
                is_nested = True
        self.is_nested_encoder = bool(is_nested)

        # Check required weights
        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(
                f"MinimalSVDAttention missing required weights for projections: {missing}. Proceeding with initialized weights."
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"MinimalSVDAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"MinimalSVDAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
                )

        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

        # Copy into factorized projection modules
        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target.decompose_from_weight(w_oriented, b, self.energy)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None, self.energy)

                target.freeze_base_weights()

        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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        if self.q_proj.down.weight.device != device or self.q_proj.down.weight.dtype != dtype:
            self.to(device=device, dtype=dtype)

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

        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, N_kv, self.num_kv_heads, D).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, N_kv, 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 and not is_cross:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = 1.0 / math.sqrt(D)
        # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

        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 = (
                getattr(self, "is_decoder", False) and not is_cross and N > 1
                if "is_cross" in dir()
                else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
            )

            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None

        except Exception:
            pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
                elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                    scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
                else:
                    scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask, -10000.0)
                else:
                    scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(dtype)
        attn_weights = self.dropout(attn_weights)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn_weights = attn_weights * head_mask

        context = torch.matmul(attn_weights, V)
        context = context.transpose(1, 2).contiguous().view(B, N, H * D)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return output, (attn_weights if output_attentions else None)

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.tensor(0.0, device=self.out_proj.up.weight.device)

MinimalSVD: SVD-Factorized Low-Rank Projection Attention. Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.

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

Ancestors

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

Subclasses

• NaiveSVDAttention

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.

Instance variables

prop key : FactorisedLinear

Expand source code

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

prop query : FactorisedLinear

Expand source code

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

prop value : FactorisedLinear

Expand source code

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

Methods

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

Expand source code

def copy_weights(self, src: Any) -> None:
    from neuralops.core.adapters import UniversalAdapter

    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 as e:
            logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

    # 2. GPT-2 Fused c_attn / c_proj 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_attn_mod = src.c_attn
        base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
        weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
        if weight is not None:
            weight = weight.data.clone()
            bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            bias_data = bias.data.clone() if bias is not None else None
            p_cls = getattr(base_c, "__class__", type(base_c)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(base_c, "nf")
                or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                or hasattr(c_attn_mod, "nf")
                or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
            ):
                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:
        c_proj_mod = src.c_proj
        base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
        w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
        if w_o is not None:
            w_o = w_o.data.clone()
            p_cls = getattr(base_o, "__class__", type(base_o)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(base_o, "nf")
                or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                or hasattr(c_proj_mod, "nf")
                or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(base_o, "bias", getattr(c_proj_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"],
        "k": ["k_proj", "k_lin", "key"],
        "v": ["v_proj", "v_lin", "value"],
        "o": ["o_proj", "out_proj", "out_lin", "dense"],
    }
    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 = src.__class__.__name__
    has_out_proj = any(
        hasattr(src, attr)
        for attr in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    is_nested = False
    if (
        hasattr(src, "self")
        and hasattr(src.self, "query")
        or (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
            )
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
            or (
                ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                and not has_out_proj
            )
        )
    ):
        is_nested = True
    if hasattr(src, "config"):
        if (
            getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
            and not has_out_proj
        ):
            is_nested = True
    self.is_nested_encoder = bool(is_nested)

    # Check required weights
    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(
            f"MinimalSVDAttention missing required weights for projections: {missing}. Proceeding with initialized weights."
        )

    # Frobenius norm validation
    for proj_key in ["q", "k", "v"]:
        w = extracted[proj_key][0]
        if w is not None:
            norm = torch.linalg.norm(w.float())
            if norm < 1e-4:
                logger.warning(
                    f"MinimalSVDAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                )

    if extracted["o"][0] is not None:
        norm_o = torch.linalg.norm(extracted["o"][0].float())
        if norm_o < 1e-4:
            logger.warning(
                f"MinimalSVDAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
            )

    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

    # Copy into factorized projection modules
    with torch.no_grad():
        for proj_key, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                w_oriented = _orient(w, in_dim)
                target.decompose_from_weight(w_oriented, b, self.energy)
            elif proj_key == "o" and self.is_nested_encoder:
                eye_w = torch.eye(
                    self.all_head_size,
                    self.all_head_size,
                    device=self.out_proj.up.weight.device,
                    dtype=self.out_proj.up.weight.dtype,
                )
                target.decompose_from_weight(eye_w, None, self.energy)

            target.freeze_base_weights()

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

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    H, D = self.num_heads, self.head_dim
    device = hidden_states.device
    dtype = hidden_states.dtype

    if self.q_proj.down.weight.device != device or self.q_proj.down.weight.dtype != dtype:
        self.to(device=device, dtype=dtype)

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

    Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
    K = self.k_proj(kv_input).view(B, N_kv, self.num_kv_heads, D).transpose(1, 2).contiguous()
    V = self.v_proj(kv_input).view(B, N_kv, 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 and not is_cross:
        cos, sin = position_embeddings
        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

    if self.num_kv_heads < self.num_heads:
        K = repeat_kv(K, self.kv_repeat)
        V = repeat_kv(V, self.kv_repeat)

    scale = 1.0 / math.sqrt(D)
    # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

    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 = (
            getattr(self, "is_decoder", False) and not is_cross and N > 1
            if "is_cross" in dir()
            else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
        )

        _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

        return (
            self.out_proj(_sdpa_out)
            if not getattr(self, "is_nested_encoder", False)
            else _sdpa_out
        ), None

    except Exception:
        pass

    scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

    if self.is_decoder and not is_cross and N > 1:
        causal_mask = (
            torch.triu(
                torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                diagonal=1,
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        scores = scores + causal_mask

    eff_mask = encoder_attention_mask if is_cross else attention_mask
    if eff_mask is not None:
        if eff_mask.dim() == 2:
            if eff_mask.dtype == torch.bool:
                scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
            elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
            else:
                scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
        elif eff_mask.dim() == 3:
            scores = scores + eff_mask.unsqueeze(1).float()
        elif eff_mask.dim() == 4:
            if eff_mask.dtype == torch.bool:
                scores = scores.masked_fill(~eff_mask, -10000.0)
            else:
                scores = scores + eff_mask.float()

    scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
    attn_weights = F.softmax(scores, dim=-1).to(dtype)
    attn_weights = self.dropout(attn_weights)

    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)
        elif head_mask.dim() == 3:
            head_mask = head_mask.unsqueeze(-1)
        attn_weights = attn_weights * head_mask

    context = torch.matmul(attn_weights, V)
    context = context.transpose(1, 2).contiguous().view(B, N, H * D)
    context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

    return output, (attn_weights if output_attentions else None)

Must return (output, attention_weights_or_None).

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.tensor(0.0, device=self.out_proj.up.weight.device)

class MinimalSVDAttentionV2 (config: Any | None = None,
model_config: Any | None = None,
energy: float = 0.95,
dropout: float = 0.1,
pretrained_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("minimal_svd", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_aurora", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_prism", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("minimal_svd_v2", supports_causal=True, supports_gqa=True)
class MinimalSVDAttention(BaseAttention):
    r"""
    MinimalSVD: SVD-Factorized Low-Rank Projection Attention.
    Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        energy: float = 0.95,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        rank: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = model_config if model_config is not None else config
        actual_rank = getattr(
            cfg,
            "svd_rank",
            rank if rank is not None else getattr(cfg, "rank", 64),
        )
        super().__init__(cfg, actual_rank)
        self.num_heads: int = getattr(
            cfg,
            "num_attention_heads",
            getattr(cfg, "num_heads", kwargs.get("num_attention_heads", 4)),
        )
        self.hidden_size: int = getattr(cfg, "hidden_size", kwargs.get("hidden_size", 768))
        cfg_head_dim = (
            getattr(cfg, "head_dim", None)
            if cfg
            else (getattr(model_config, "head_dim", None) if model_config else None)
        )
        self.head_dim: int = int(
            cfg_head_dim
            or kwargs.get("head_dim")
            or (
                self.hidden_size // self.num_heads
                if self.num_heads > 0
                else (self.hidden_size or 64)
            )
        )
        self.all_head_size: int = self.num_heads * self.head_dim

        cfg_kv = (
            getattr(cfg, "num_key_value_heads", getattr(cfg, "num_kv_heads", None)) if cfg else None
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        self.rank: int = actual_rank
        self.energy: float = energy
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = (
                getattr(cfg, "is_decoder", False)
                or getattr(cfg, "causal", False)
                or getattr(cfg, "model_type", "").lower() in ("gpt2", "qwen", "qwen2", "llama")
                or any(
                    str(
                        getattr(
                            (
                                cfg
                                if "cfg" in locals()
                                else (config if "config" in locals() else None)
                            ),
                            "model_type",
                            "",
                        )
                    )
                    .lower()
                    .startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                if cfg
                else False
            )

        model_type = getattr(cfg, "model_type", "").lower() if cfg else ""
        self.is_nested_encoder: bool = getattr(cfg, "is_nested_encoder", False) or model_type in (
            "roberta",
            "bert",
            "xlm-roberta",
        )

        self.q_proj = FactorisedLinear(
            self.hidden_size,
            self.all_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.k_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.v_proj = FactorisedLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            self.hidden_size,
            rank=self.rank,
            bias=True,
            energy=energy,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", dropout))
            if cfg
            else dropout
        )
        self.dropout: nn.Dropout = nn.Dropout(dropout_p)

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

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

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

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

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug(f"UniversalAdapter extraction failed for {proj}: {e}")

        # 2. GPT-2 Fused c_attn / c_proj 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_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                p_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
                ):
                    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:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_o, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(c_proj_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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
                or (
                    ("roberta" in cls_name.lower() or "bert" in cls_name.lower())
                    and not has_out_proj
                )
            )
        ):
            is_nested = True
        if hasattr(src, "config"):
            if (
                getattr(src.config, "model_type", "").lower() in ("roberta", "bert", "xlm-roberta")
                and not has_out_proj
            ):
                is_nested = True
        self.is_nested_encoder = bool(is_nested)

        # Check required weights
        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(
                f"MinimalSVDAttention missing required weights for projections: {missing}. Proceeding with initialized weights."
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"MinimalSVDAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"MinimalSVDAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
                )

        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

        # Copy into factorized projection modules
        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target.decompose_from_weight(w_oriented, b, self.energy)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=self.out_proj.up.weight.device,
                        dtype=self.out_proj.up.weight.dtype,
                    )
                    target.decompose_from_weight(eye_w, None, self.energy)

                target.freeze_base_weights()

        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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        if self.q_proj.down.weight.device != device or self.q_proj.down.weight.dtype != dtype:
            self.to(device=device, dtype=dtype)

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

        Q = self.q_proj(hidden_states).view(B, N, H, D).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, N_kv, self.num_kv_heads, D).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, N_kv, 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 and not is_cross:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = 1.0 / math.sqrt(D)
        # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

        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 = (
                getattr(self, "is_decoder", False) and not is_cross and N > 1
                if "is_cross" in dir()
                else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
            )

            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None

        except Exception:
            pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
                elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                    scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
                else:
                    scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask, -10000.0)
                else:
                    scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(dtype)
        attn_weights = self.dropout(attn_weights)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn_weights = attn_weights * head_mask

        context = torch.matmul(attn_weights, V)
        context = context.transpose(1, 2).contiguous().view(B, N, H * D)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

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

        return output, (attn_weights if output_attentions else None)

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.tensor(0.0, device=self.out_proj.up.weight.device)

MinimalSVD: SVD-Factorized Low-Rank Projection Attention. Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.

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

Ancestors

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

Subclasses

• NaiveSVDAttention

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.

Instance variables

prop key : FactorisedLinear

Expand source code

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

prop query : FactorisedLinear

Expand source code

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

prop value : FactorisedLinear

Expand source code

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

Methods

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.tensor(0.0, device=self.out_proj.up.weight.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class FASA (config: Any | None = None,
rank: int | None = 32,
fasa_keep_ratio: float = 0.5,
pretrained_attn: torch.nn.modules.module.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("fasa", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("fasa_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("mixattention_moska", supports_causal=True, supports_gqa=True)
class MixAttentionMoSKA(BaseAttention):
    """
    MixAttention-MoSKA (FASA): frequency/variance-aware sparse attention.

    For non-causal self-attention and encoder-style models, we prune the K/V
    set down to the top `fasa_keep_ratio` tokens by variance importance,
    always preserving the [CLS] token at index 0.

    For causal decoders, cross-attention, and incremental decoding (when
    `past_key_value` is present), we bypass pruning entirely and run exact
    dense attention via SDPA / manual softmax.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 32,
        fasa_keep_ratio: float = 0.5,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        # FASA-FIX-04: config None-safe
        if config is None:

            class _Cfg:
                pass

            config = _Cfg()
            config.hidden_size = kwargs.get("hidden_size", 768)
            config.num_attention_heads = kwargs.get(
                "num_attention_heads", kwargs.get("num_heads", 12)
            )
            config.head_dim = kwargs.get("head_dim", 64)
            config.model_type = kwargs.get("model_type", "")
            config.is_decoder = kwargs.get("is_decoder", False)
            config.causal = kwargs.get("is_causal", False)

        try:
            super().__init__(config, rank if rank is not None else 32)
        except TypeError:
            super().__init__(config)

        self.dim = int(
            kwargs.get("hidden_size") or getattr(config, "hidden_size", getattr(config, "dim", 768))
        )
        self.hidden_size = self.dim

        n_heads = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or getattr(config, "num_attention_heads", getattr(config, "n_heads", 0))
        )
        # FASA-FIX-11
        if not n_heads or int(n_heads) <= 0:
            n_heads = max(1, self.dim // 64)
        self.num_heads = int(n_heads)

        cfg_head_dim = getattr(config, "head_dim", None)
        self.head_dim = int(
            cfg_head_dim or kwargs.get("head_dim") or (self.dim // max(1, self.num_heads))
        )
        self.all_head_size = self.num_heads * self.head_dim

        self.fasa_keep_ratio = float(fasa_keep_ratio)
        self.protect_first_token = bool(kwargs.get("protect_first_token", True))
        self.importance_method = str(kwargs.get("importance_method", "var"))

        cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = self.num_heads // self.num_kv_heads
        self.kv_head_size = self.num_kv_heads * self.head_dim

        self.lora_rank = int(lora_rank)
        self.lora_alpha = float(lora_alpha)

        # ---- decoder detection ----
        _mt = str(getattr(config, "model_type", "") or "").lower()
        _dec_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = bool(
                getattr(config, "is_decoder", False)
                or getattr(config, "causal", False)
                or any(_mt.startswith(p) for p in _dec_prefixes)
                or kwargs.get("is_causal", False)
            )
            # FASA-FIX-01: also detect GPT-2 by source attention
            if pretrained_attn is not None:
                if hasattr(pretrained_attn, "c_attn"):
                    self.is_decoder = True
                if hasattr(pretrained_attn, "is_decoder") and pretrained_attn.is_decoder:
                    self.is_decoder = True

        self.is_nested_encoder = False

        self.W_q = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.W_k = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.W_v = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
        )
        self.dropout = nn.Dropout(dropout_p)

        # FASA-FIX-09: storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

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

    # ---- aliases ----
    @property
    def query(self):
        return self.W_q

    @property
    def key(self):
        return self.W_k

    @property
    def value(self):
        return self.W_v

    @property
    def q_proj(self):
        return self.W_q

    @property
    def k_proj(self):
        return self.W_k

    @property
    def v_proj(self):
        return self.W_v

    @property
    def q_lin(self):
        return self.W_q

    @property
    def k_lin(self):
        return self.W_k

    @property
    def v_lin(self):
        return self.W_v

    @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

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        # Track original storage
        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. UniversalAdapter
        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
        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 = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split],
                    bias_data[:split] if bias_data is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split],
                    bias_data[split : 2 * split] if bias_data is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :],
                    bias_data[2 * split :] if bias_data is not None else None,
                )

        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. BERT / RoBERTa nested (FASA-FIX-02: extract .data, not module)
        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
            for key, sub_name in (("q", "query"), ("k", "key"), ("v", "value")):
                if extracted[key][0] is None:
                    m = getattr(s, sub_name, None)
                    if m is not None and hasattr(m, "weight"):
                        bias = getattr(m, "bias", None)
                        extracted[key] = (
                            m.weight.data,
                            bias.data if bias is not None else None,
                        )

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            d = getattr(src.output.dense, "base_layer", src.output.dense)
            if hasattr(d, "weight"):
                bias = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data,
                    bias.data if bias is not None else None,
                )

        # 4. Flat modules
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(src, name):
                        m = getattr(src, name)
                        if hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[proj] = (
                                m.weight.data,
                                b.data if b is not None else None,
                            )
                            break

        # ---- nested encoder detection ----
        cls_name = getattr(src, "__class__", type(src)).__name__
        has_out = any(hasattr(src, a) for a in ("dense", "out_proj", "out_lin", "c_proj", "o_proj"))
        is_nested = False
        if (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
                "DebertaSelfAttention",
            )
            or hasattr(src, "self")
            and hasattr(src, "output")
            or hasattr(src, "query")
            and not has_out
        ):
            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("MixAttentionMoSKA copy_weights: missing %s", missing)

        # Frobenius norm checks
        for k_ in ("q", "k", "v"):
            w = extracted[k_][0]
            if w is not None:
                n = torch.linalg.norm(w.float()).item()
                if n < 1e-4:
                    logger.warning("FASA: %s weight near-zero norm (%.2e)", k_, n)
        if extracted["o"][0] is not None:
            n = torch.linalg.norm(extracted["o"][0].float()).item()
            if n < 1e-4:
                logger.warning("FASA: o weight near-zero norm (%.2e)", n)

        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

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.W_q),
                ("k", self.W_k),
                ("v", self.W_v),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.dim
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape

                    if proj_key in ("k", "v"):
                        if src_out == out_d:
                            target_w.data.copy_(w_oriented)
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b)
                        elif src_out > out_d and src_out % out_d == 0:
                            rep = src_out // out_d
                            target_w.data.copy_(w_oriented.view(out_d, rep, -1).mean(dim=1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.view(out_d, rep).mean(dim=1))
                        elif src_out < out_d and out_d % src_out == 0:
                            rep = out_d // src_out
                            target_w.data.copy_(w_oriented.repeat(rep, 1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.repeat(rep))
                        else:
                            m0 = min(out_d, src_out)
                            m1 = min(in_d, src_in)
                            target_w.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                            if b is not None and target.bias is not None:
                                mb = min(target.bias.shape[0], b.shape[0])
                                target.bias.data[:mb].copy_(b[:mb])
                    else:
                        m0 = min(out_d, src_out)
                        m1 = min(in_d, src_in)
                        target_w.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                        if b is not None and target.bias is not None:
                            mb = min(target.bias.shape[0], b.shape[0])
                            target.bias.data[:mb].copy_(b[:mb])
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                target.freeze_base_weights()

        # Track compressed storage
        try:
            for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # FASA-FIX-06: don't overwrite if src lacks is_decoder
        if hasattr(src, "is_decoder") and src.is_decoder:
            self.is_decoder = True

    # ------------------------------------------------------------------
    # FASA pruning
    # ------------------------------------------------------------------
    def _fasa_prune(
        self,
        k: torch.Tensor,
        v: torch.Tensor,
        mask_2d: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
        B, H, L, D_h = k.shape
        if L <= 8:
            return k, v, mask_2d, None

        method = getattr(self, "importance_method", "var")
        if method == "l2":
            importance = torch.norm(k.float(), p=2, dim=-1).mean(dim=1)  # (B, L)
        else:
            importance = k.float().var(dim=-1).mean(dim=1)  # (B, L)

        if mask_2d is not None:
            importance = importance.masked_fill(~mask_2d, -1e4)

        # FASA-FIX-03: clone before in-place modification
        importance = importance.clone()
        if self.protect_first_token:
            importance[:, 0] = 1e9

        keep = max(8, min(L, int(math.ceil(L * self.fasa_keep_ratio))))
        _, idx = torch.topk(importance, keep, dim=-1)
        idx, _ = torch.sort(idx, dim=-1)

        idx_exp = idx.unsqueeze(1).unsqueeze(-1).expand(-1, H, -1, D_h)
        k_pruned = k.gather(2, idx_exp)
        v_pruned = v.gather(2, idx_exp)

        pruned_mask = None
        if mask_2d is not None:
            pruned_mask = mask_2d.gather(1, idx)

        return k_pruned, v_pruned, pruned_mask, idx

    # ------------------------------------------------------------------
    # Mask helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        L: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """Return a bool (B, L) mask where True = valid."""
        if attention_mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape[-1] == L:
                return m
        elif attention_mask.dim() == 3:
            m = _to_bool(attention_mask).squeeze(1)
            if m.shape[-1] == L:
                return m
        elif attention_mask.dim() == 4:
            m = _to_bool(attention_mask).any(dim=-2).squeeze(1)
            if m.shape[-1] == L:
                return m

        # Fallback: try the framework helper
        try:
            m = get_token_valid_mask(attention_mask, B, L, device)
            if isinstance(m, torch.Tensor) and m.shape == (B, L):
                return m.bool()
        except Exception:
            pass
        return None

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ):
        # FASA-FIX-07: fold `mask` kwarg into attention_mask
        mask = kwargs.get("mask")
        if attention_mask is None and mask is not None:
            attention_mask = mask

        x = hidden_states
        B, L, D = x.shape
        H = self.num_heads
        device = x.device
        dtype = x.dtype

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else x
        L_kv = kv_input.shape[1]

        q = self.W_q(x).view(B, L, H, self.head_dim).permute(0, 2, 1, 3)
        k = self.W_k(kv_input).view(B, L_kv, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
        v = self.W_v(kv_input).view(B, L_kv, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)

        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
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        if self.num_kv_heads < self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        scale = 1.0 / math.sqrt(self.head_dim)
        is_causal = bool(self.is_decoder and not is_cross and L > 1)

        # FASA-FIX-01: disable pruning when:
        #   - we are in a causal decoder,
        #   - we are doing cross-attention,
        #   - a KV-cache is present (incremental decoding),
        #   - hidden_states and kv_input have different lengths (mixed mode).
        use_pruning = not is_causal and not is_cross and past_key_value is None and L_kv == L

        if not use_pruning:
            # -------- Dense path (SDPA / manual) --------
            if eff_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
                    _sdpa_out = F.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(dtype)
                        .transpose(1, 2)
                        .contiguous()
                        .view(B, L, self.all_head_size)
                    )
                    _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                    if self.is_nested_encoder:
                        return _sdpa_out, None
                    return self.out_proj(_sdpa_out), None
                except Exception:
                    pass

            scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
            if is_causal:
                causal = (
                    torch.triu(
                        torch.full((L, L_kv), float("-inf"), device=device, dtype=torch.float32),
                        diagonal=1,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal

            valid = self._normalize_valid_mask(eff_mask, B, L_kv, device)
            if valid is not None:
                scores = scores.masked_fill(~valid[:, None, None, :], -1e4)

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            attn_weights = F.softmax(scores, dim=-1).to(dtype)
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            attn_weights = self.dropout(attn_weights)

            # FASA-FIX-10: unified head_mask handling
            if head_mask is not None:
                hm = head_mask.to(attn_weights.dtype)
                if hm.dim() == 1:
                    hm = hm.view(1, -1, 1, 1)
                elif hm.dim() == 2:
                    hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
                elif hm.dim() == 3:
                    hm = hm.unsqueeze(1)
                attn_weights = attn_weights * hm

            out = (
                torch.matmul(attn_weights, v)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.all_head_size)
            )
            final_attn_weights = attn_weights if output_attentions else None
        else:
            # -------- FASA pruning path --------
            valid = self._normalize_valid_mask(eff_mask, B, L_kv, device)
            k_p, v_p, p_mask, kept_idx = self._fasa_prune(k, v, valid)

            attn_scores = (
                torch.matmul(q.float(), k_p.float().transpose(-2, -1).contiguous()) * scale
            )

            if p_mask is not None:
                attn_scores = attn_scores.masked_fill(~p_mask[:, None, None, :], -1e4)

            attn_scores = torch.nan_to_num(attn_scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            attn_weights = F.softmax(attn_scores, dim=-1).to(dtype)
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            attn_weights = self.dropout(attn_weights)

            if head_mask is not None:
                hm = head_mask.to(attn_weights.dtype)
                if hm.dim() == 1:
                    hm = hm.view(1, -1, 1, 1)
                elif hm.dim() == 2:
                    hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
                elif hm.dim() == 3:
                    hm = hm.unsqueeze(1)
                attn_weights = attn_weights * hm

            out = (
                torch.matmul(attn_weights, v_p)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.all_head_size)
            )

            if output_attentions and kept_idx is not None:
                full_weights = torch.zeros(B, H, L, L_kv, device=device, dtype=dtype)
                kept_exp = kept_idx.unsqueeze(1).unsqueeze(2).expand(-1, H, L, -1)
                final_attn_weights = full_weights.scatter(3, kept_exp, attn_weights)
            elif output_attentions:
                final_attn_weights = attn_weights
            else:
                final_attn_weights = None

        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)

        return output, final_attn_weights

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if p is None:
                return 0
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("W_q", self.W_q),
            ("W_k", self.W_k),
            ("W_v", self.W_v),
            ("out_proj", self.out_proj),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["total"] = sum(out.values())
        return out

MixAttention-MoSKA (FASA): frequency/variance-aware sparse attention.

For non-causal self-attention and encoder-style models, we prune the K/V set down to the top fasa_keep_ratio tokens by variance importance, always preserving the [CLS] token at index 0.

For causal decoders, cross-attention, and incremental decoding (when past_key_value is present), we bypass pruning entirely and run exact dense attention via SDPA / manual softmax.

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.

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.W_k

prop k_proj

Expand source code

@property
def k_proj(self):
    return self.W_k

prop key

Expand source code

@property
def key(self):
    return self.W_k

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.W_q

prop q_proj

Expand source code

@property
def q_proj(self):
    return self.W_q

prop query

Expand source code

@property
def query(self):
    return self.W_q

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.W_v

prop v_proj

Expand source code

@property
def v_proj(self):
    return self.W_v

prop value

Expand source code

@property
def value(self):
    return self.W_v

Methods

def copy_weights(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    from neuralops.core.adapters import UniversalAdapter

    # Track original storage
    try:
        for p in src.parameters():
            self._orig_storage_bytes += p.numel() * p.element_size()
    except Exception:
        pass

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

    # 1. UniversalAdapter
    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
    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 = weight.shape[0] // 3
            extracted["q"] = (
                weight[:split],
                bias_data[:split] if bias_data is not None else None,
            )
            extracted["k"] = (
                weight[split : 2 * split],
                bias_data[split : 2 * split] if bias_data is not None else None,
            )
            extracted["v"] = (
                weight[2 * split :],
                bias_data[2 * split :] if bias_data is not None else None,
            )

    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. BERT / RoBERTa nested (FASA-FIX-02: extract .data, not module)
    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
        for key, sub_name in (("q", "query"), ("k", "key"), ("v", "value")):
            if extracted[key][0] is None:
                m = getattr(s, sub_name, None)
                if m is not None and hasattr(m, "weight"):
                    bias = getattr(m, "bias", None)
                    extracted[key] = (
                        m.weight.data,
                        bias.data if bias is not None else None,
                    )

    if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
        d = getattr(src.output.dense, "base_layer", src.output.dense)
        if hasattr(d, "weight"):
            bias = getattr(d, "bias", None)
            extracted["o"] = (
                d.weight.data,
                bias.data if bias is not None else None,
            )

    # 4. Flat modules
    named_map = {
        "q": ("q_proj", "q_lin", "query"),
        "k": ("k_proj", "k_lin", "key"),
        "v": ("v_proj", "v_lin", "value"),
        "o": ("o_proj", "out_proj", "out_lin", "dense"),
    }
    for proj, names in named_map.items():
        if extracted[proj][0] is None:
            for name in names:
                if hasattr(src, name):
                    m = getattr(src, name)
                    if hasattr(m, "weight"):
                        b = getattr(m, "bias", None)
                        extracted[proj] = (
                            m.weight.data,
                            b.data if b is not None else None,
                        )
                        break

    # ---- nested encoder detection ----
    cls_name = getattr(src, "__class__", type(src)).__name__
    has_out = any(hasattr(src, a) for a in ("dense", "out_proj", "out_lin", "c_proj", "o_proj"))
    is_nested = False
    if (
        cls_name
        in (
            "BertSelfAttention",
            "RobertaSelfAttention",
            "XLMRobertaSelfAttention",
            "ElectraSelfAttention",
            "DebertaSelfAttention",
        )
        or hasattr(src, "self")
        and hasattr(src, "output")
        or hasattr(src, "query")
        and not has_out
    ):
        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("MixAttentionMoSKA copy_weights: missing %s", missing)

    # Frobenius norm checks
    for k_ in ("q", "k", "v"):
        w = extracted[k_][0]
        if w is not None:
            n = torch.linalg.norm(w.float()).item()
            if n < 1e-4:
                logger.warning("FASA: %s weight near-zero norm (%.2e)", k_, n)
    if extracted["o"][0] is not None:
        n = torch.linalg.norm(extracted["o"][0].float()).item()
        if n < 1e-4:
            logger.warning("FASA: o weight near-zero norm (%.2e)", n)

    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

    with torch.no_grad():
        for proj_key, target in (
            ("q", self.W_q),
            ("k", self.W_k),
            ("v", self.W_v),
            ("o", self.out_proj),
        ):
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.dim
                w_oriented = _orient(w, in_dim)
                target_w = target.weight
                out_d, in_d = target_w.shape
                src_out, src_in = w_oriented.shape

                if proj_key in ("k", "v"):
                    if src_out == out_d:
                        target_w.data.copy_(w_oriented)
                        if b is not None and target.bias is not None:
                            target.bias.data.copy_(b)
                    elif src_out > out_d and src_out % out_d == 0:
                        rep = src_out // out_d
                        target_w.data.copy_(w_oriented.view(out_d, rep, -1).mean(dim=1))
                        if b is not None and target.bias is not None:
                            target.bias.data.copy_(b.view(out_d, rep).mean(dim=1))
                    elif src_out < out_d and out_d % src_out == 0:
                        rep = out_d // src_out
                        target_w.data.copy_(w_oriented.repeat(rep, 1))
                        if b is not None and target.bias is not None:
                            target.bias.data.copy_(b.repeat(rep))
                    else:
                        m0 = min(out_d, src_out)
                        m1 = min(in_d, src_in)
                        target_w.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                        if b is not None and target.bias is not None:
                            mb = min(target.bias.shape[0], b.shape[0])
                            target.bias.data[:mb].copy_(b[:mb])
                else:
                    m0 = min(out_d, src_out)
                    m1 = min(in_d, src_in)
                    target_w.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                    if b is not None and target.bias is not None:
                        mb = min(target.bias.shape[0], b.shape[0])
                        target.bias.data[:mb].copy_(b[:mb])
            elif proj_key == "o" and self.is_nested_encoder:
                min_d = min(target.weight.shape[0], target.weight.shape[1])
                target.weight.data.zero_()
                target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                if target.bias is not None:
                    target.bias.data.zero_()

            target.freeze_base_weights()

    # Track compressed storage
    try:
        for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
            for p in m.parameters():
                self._comp_storage_bytes += p.numel() * p.element_size()
    except Exception:
        pass

    # FASA-FIX-06: don't overwrite if src lacks is_decoder
    if hasattr(src, "is_decoder") and src.is_decoder:
        self.is_decoder = True

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,
*args: Any,
**kwargs: Any) ‑> Callable[..., typing.Any]

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,
    *args: Any,
    **kwargs: Any,
):
    # FASA-FIX-07: fold `mask` kwarg into attention_mask
    mask = kwargs.get("mask")
    if attention_mask is None and mask is not None:
        attention_mask = mask

    x = hidden_states
    B, L, D = x.shape
    H = self.num_heads
    device = x.device
    dtype = x.dtype

    is_cross = encoder_hidden_states is not None
    kv_input = encoder_hidden_states if is_cross else x
    L_kv = kv_input.shape[1]

    q = self.W_q(x).view(B, L, H, self.head_dim).permute(0, 2, 1, 3)
    k = self.W_k(kv_input).view(B, L_kv, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
    v = self.W_v(kv_input).view(B, L_kv, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)

    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
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    if self.num_kv_heads < self.num_heads:
        k = repeat_kv(k, self.kv_repeat)
        v = repeat_kv(v, self.kv_repeat)

    eff_mask = encoder_attention_mask if is_cross else attention_mask
    scale = 1.0 / math.sqrt(self.head_dim)
    is_causal = bool(self.is_decoder and not is_cross and L > 1)

    # FASA-FIX-01: disable pruning when:
    #   - we are in a causal decoder,
    #   - we are doing cross-attention,
    #   - a KV-cache is present (incremental decoding),
    #   - hidden_states and kv_input have different lengths (mixed mode).
    use_pruning = not is_causal and not is_cross and past_key_value is None and L_kv == L

    if not use_pruning:
        # -------- Dense path (SDPA / manual) --------
        if eff_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
                _sdpa_out = F.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(dtype)
                    .transpose(1, 2)
                    .contiguous()
                    .view(B, L, self.all_head_size)
                )
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                if self.is_nested_encoder:
                    return _sdpa_out, None
                return self.out_proj(_sdpa_out), None
            except Exception:
                pass

        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
        if is_causal:
            causal = (
                torch.triu(
                    torch.full((L, L_kv), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        valid = self._normalize_valid_mask(eff_mask, B, L_kv, device)
        if valid is not None:
            scores = scores.masked_fill(~valid[:, None, None, :], -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(dtype)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights)

        # FASA-FIX-10: unified head_mask handling
        if head_mask is not None:
            hm = head_mask.to(attn_weights.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            attn_weights = attn_weights * hm

        out = (
            torch.matmul(attn_weights, v)
            .transpose(1, 2)
            .contiguous()
            .view(B, L, self.all_head_size)
        )
        final_attn_weights = attn_weights if output_attentions else None
    else:
        # -------- FASA pruning path --------
        valid = self._normalize_valid_mask(eff_mask, B, L_kv, device)
        k_p, v_p, p_mask, kept_idx = self._fasa_prune(k, v, valid)

        attn_scores = (
            torch.matmul(q.float(), k_p.float().transpose(-2, -1).contiguous()) * scale
        )

        if p_mask is not None:
            attn_scores = attn_scores.masked_fill(~p_mask[:, None, None, :], -1e4)

        attn_scores = torch.nan_to_num(attn_scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(attn_scores, dim=-1).to(dtype)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights)

        if head_mask is not None:
            hm = head_mask.to(attn_weights.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            attn_weights = attn_weights * hm

        out = (
            torch.matmul(attn_weights, v_p)
            .transpose(1, 2)
            .contiguous()
            .view(B, L, self.all_head_size)
        )

        if output_attentions and kept_idx is not None:
            full_weights = torch.zeros(B, H, L, L_kv, device=device, dtype=dtype)
            kept_exp = kept_idx.unsqueeze(1).unsqueeze(2).expand(-1, H, L, -1)
            final_attn_weights = full_weights.scatter(3, kept_exp, attn_weights)
        elif output_attentions:
            final_attn_weights = attn_weights
        else:
            final_attn_weights = None

    out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

    if self.is_nested_encoder:
        output = out
    else:
        output = self.out_proj(out)

    return output, final_attn_weights

Must return (output, attention_weights_or_None).

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if p is None:
            return 0
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("W_q", self.W_q),
        ("W_k", self.W_k),
        ("W_v", self.W_v),
        ("out_proj", self.out_proj),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["total"] = sum(out.values())
    return out

class MixAttentionMoSKA (config: Any | None = None,
rank: int | None = 32,
fasa_keep_ratio: float = 0.5,
pretrained_attn: torch.nn.modules.module.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("fasa", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("fasa_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("mixattention_moska", supports_causal=True, supports_gqa=True)
class MixAttentionMoSKA(BaseAttention):
    """
    MixAttention-MoSKA (FASA): frequency/variance-aware sparse attention.

    For non-causal self-attention and encoder-style models, we prune the K/V
    set down to the top `fasa_keep_ratio` tokens by variance importance,
    always preserving the [CLS] token at index 0.

    For causal decoders, cross-attention, and incremental decoding (when
    `past_key_value` is present), we bypass pruning entirely and run exact
    dense attention via SDPA / manual softmax.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 32,
        fasa_keep_ratio: float = 0.5,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        # FASA-FIX-04: config None-safe
        if config is None:

            class _Cfg:
                pass

            config = _Cfg()
            config.hidden_size = kwargs.get("hidden_size", 768)
            config.num_attention_heads = kwargs.get(
                "num_attention_heads", kwargs.get("num_heads", 12)
            )
            config.head_dim = kwargs.get("head_dim", 64)
            config.model_type = kwargs.get("model_type", "")
            config.is_decoder = kwargs.get("is_decoder", False)
            config.causal = kwargs.get("is_causal", False)

        try:
            super().__init__(config, rank if rank is not None else 32)
        except TypeError:
            super().__init__(config)

        self.dim = int(
            kwargs.get("hidden_size") or getattr(config, "hidden_size", getattr(config, "dim", 768))
        )
        self.hidden_size = self.dim

        n_heads = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or getattr(config, "num_attention_heads", getattr(config, "n_heads", 0))
        )
        # FASA-FIX-11
        if not n_heads or int(n_heads) <= 0:
            n_heads = max(1, self.dim // 64)
        self.num_heads = int(n_heads)

        cfg_head_dim = getattr(config, "head_dim", None)
        self.head_dim = int(
            cfg_head_dim or kwargs.get("head_dim") or (self.dim // max(1, self.num_heads))
        )
        self.all_head_size = self.num_heads * self.head_dim

        self.fasa_keep_ratio = float(fasa_keep_ratio)
        self.protect_first_token = bool(kwargs.get("protect_first_token", True))
        self.importance_method = str(kwargs.get("importance_method", "var"))

        cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = self.num_heads // self.num_kv_heads
        self.kv_head_size = self.num_kv_heads * self.head_dim

        self.lora_rank = int(lora_rank)
        self.lora_alpha = float(lora_alpha)

        # ---- decoder detection ----
        _mt = str(getattr(config, "model_type", "") or "").lower()
        _dec_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = bool(
                getattr(config, "is_decoder", False)
                or getattr(config, "causal", False)
                or any(_mt.startswith(p) for p in _dec_prefixes)
                or kwargs.get("is_causal", False)
            )
            # FASA-FIX-01: also detect GPT-2 by source attention
            if pretrained_attn is not None:
                if hasattr(pretrained_attn, "c_attn"):
                    self.is_decoder = True
                if hasattr(pretrained_attn, "is_decoder") and pretrained_attn.is_decoder:
                    self.is_decoder = True

        self.is_nested_encoder = False

        self.W_q = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.W_k = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.W_v = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
        )
        self.dropout = nn.Dropout(dropout_p)

        # FASA-FIX-09: storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    # ---- aliases ----
    @property
    def query(self):
        return self.W_q

    @property
    def key(self):
        return self.W_k

    @property
    def value(self):
        return self.W_v

    @property
    def q_proj(self):
        return self.W_q

    @property
    def k_proj(self):
        return self.W_k

    @property
    def v_proj(self):
        return self.W_v

    @property
    def q_lin(self):
        return self.W_q

    @property
    def k_lin(self):
        return self.W_k

    @property
    def v_lin(self):
        return self.W_v

    @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

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        # Track original storage
        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. UniversalAdapter
        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
        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 = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split],
                    bias_data[:split] if bias_data is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split],
                    bias_data[split : 2 * split] if bias_data is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :],
                    bias_data[2 * split :] if bias_data is not None else None,
                )

        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. BERT / RoBERTa nested (FASA-FIX-02: extract .data, not module)
        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
            for key, sub_name in (("q", "query"), ("k", "key"), ("v", "value")):
                if extracted[key][0] is None:
                    m = getattr(s, sub_name, None)
                    if m is not None and hasattr(m, "weight"):
                        bias = getattr(m, "bias", None)
                        extracted[key] = (
                            m.weight.data,
                            bias.data if bias is not None else None,
                        )

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            d = getattr(src.output.dense, "base_layer", src.output.dense)
            if hasattr(d, "weight"):
                bias = getattr(d, "bias", None)
                extracted["o"] = (
                    d.weight.data,
                    bias.data if bias is not None else None,
                )

        # 4. Flat modules
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(src, name):
                        m = getattr(src, name)
                        if hasattr(m, "weight"):
                            b = getattr(m, "bias", None)
                            extracted[proj] = (
                                m.weight.data,
                                b.data if b is not None else None,
                            )
                            break

        # ---- nested encoder detection ----
        cls_name = getattr(src, "__class__", type(src)).__name__
        has_out = any(hasattr(src, a) for a in ("dense", "out_proj", "out_lin", "c_proj", "o_proj"))
        is_nested = False
        if (
            cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
                "DebertaSelfAttention",
            )
            or hasattr(src, "self")
            and hasattr(src, "output")
            or hasattr(src, "query")
            and not has_out
        ):
            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("MixAttentionMoSKA copy_weights: missing %s", missing)

        # Frobenius norm checks
        for k_ in ("q", "k", "v"):
            w = extracted[k_][0]
            if w is not None:
                n = torch.linalg.norm(w.float()).item()
                if n < 1e-4:
                    logger.warning("FASA: %s weight near-zero norm (%.2e)", k_, n)
        if extracted["o"][0] is not None:
            n = torch.linalg.norm(extracted["o"][0].float()).item()
            if n < 1e-4:
                logger.warning("FASA: o weight near-zero norm (%.2e)", n)

        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

        with torch.no_grad():
            for proj_key, target in (
                ("q", self.W_q),
                ("k", self.W_k),
                ("v", self.W_v),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.dim
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape

                    if proj_key in ("k", "v"):
                        if src_out == out_d:
                            target_w.data.copy_(w_oriented)
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b)
                        elif src_out > out_d and src_out % out_d == 0:
                            rep = src_out // out_d
                            target_w.data.copy_(w_oriented.view(out_d, rep, -1).mean(dim=1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.view(out_d, rep).mean(dim=1))
                        elif src_out < out_d and out_d % src_out == 0:
                            rep = out_d // src_out
                            target_w.data.copy_(w_oriented.repeat(rep, 1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.repeat(rep))
                        else:
                            m0 = min(out_d, src_out)
                            m1 = min(in_d, src_in)
                            target_w.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                            if b is not None and target.bias is not None:
                                mb = min(target.bias.shape[0], b.shape[0])
                                target.bias.data[:mb].copy_(b[:mb])
                    else:
                        m0 = min(out_d, src_out)
                        m1 = min(in_d, src_in)
                        target_w.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                        if b is not None and target.bias is not None:
                            mb = min(target.bias.shape[0], b.shape[0])
                            target.bias.data[:mb].copy_(b[:mb])
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                target.freeze_base_weights()

        # Track compressed storage
        try:
            for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # FASA-FIX-06: don't overwrite if src lacks is_decoder
        if hasattr(src, "is_decoder") and src.is_decoder:
            self.is_decoder = True

    # ------------------------------------------------------------------
    # FASA pruning
    # ------------------------------------------------------------------
    def _fasa_prune(
        self,
        k: torch.Tensor,
        v: torch.Tensor,
        mask_2d: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
        B, H, L, D_h = k.shape
        if L <= 8:
            return k, v, mask_2d, None

        method = getattr(self, "importance_method", "var")
        if method == "l2":
            importance = torch.norm(k.float(), p=2, dim=-1).mean(dim=1)  # (B, L)
        else:
            importance = k.float().var(dim=-1).mean(dim=1)  # (B, L)

        if mask_2d is not None:
            importance = importance.masked_fill(~mask_2d, -1e4)

        # FASA-FIX-03: clone before in-place modification
        importance = importance.clone()
        if self.protect_first_token:
            importance[:, 0] = 1e9

        keep = max(8, min(L, int(math.ceil(L * self.fasa_keep_ratio))))
        _, idx = torch.topk(importance, keep, dim=-1)
        idx, _ = torch.sort(idx, dim=-1)

        idx_exp = idx.unsqueeze(1).unsqueeze(-1).expand(-1, H, -1, D_h)
        k_pruned = k.gather(2, idx_exp)
        v_pruned = v.gather(2, idx_exp)

        pruned_mask = None
        if mask_2d is not None:
            pruned_mask = mask_2d.gather(1, idx)

        return k_pruned, v_pruned, pruned_mask, idx

    # ------------------------------------------------------------------
    # Mask helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        L: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """Return a bool (B, L) mask where True = valid."""
        if attention_mask is None:
            return None

        def _to_bool(t):
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape[-1] == L:
                return m
        elif attention_mask.dim() == 3:
            m = _to_bool(attention_mask).squeeze(1)
            if m.shape[-1] == L:
                return m
        elif attention_mask.dim() == 4:
            m = _to_bool(attention_mask).any(dim=-2).squeeze(1)
            if m.shape[-1] == L:
                return m

        # Fallback: try the framework helper
        try:
            m = get_token_valid_mask(attention_mask, B, L, device)
            if isinstance(m, torch.Tensor) and m.shape == (B, L):
                return m.bool()
        except Exception:
            pass
        return None

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ):
        # FASA-FIX-07: fold `mask` kwarg into attention_mask
        mask = kwargs.get("mask")
        if attention_mask is None and mask is not None:
            attention_mask = mask

        x = hidden_states
        B, L, D = x.shape
        H = self.num_heads
        device = x.device
        dtype = x.dtype

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else x
        L_kv = kv_input.shape[1]

        q = self.W_q(x).view(B, L, H, self.head_dim).permute(0, 2, 1, 3)
        k = self.W_k(kv_input).view(B, L_kv, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
        v = self.W_v(kv_input).view(B, L_kv, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)

        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
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        if self.num_kv_heads < self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        scale = 1.0 / math.sqrt(self.head_dim)
        is_causal = bool(self.is_decoder and not is_cross and L > 1)

        # FASA-FIX-01: disable pruning when:
        #   - we are in a causal decoder,
        #   - we are doing cross-attention,
        #   - a KV-cache is present (incremental decoding),
        #   - hidden_states and kv_input have different lengths (mixed mode).
        use_pruning = not is_causal and not is_cross and past_key_value is None and L_kv == L

        if not use_pruning:
            # -------- Dense path (SDPA / manual) --------
            if eff_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
                    _sdpa_out = F.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(dtype)
                        .transpose(1, 2)
                        .contiguous()
                        .view(B, L, self.all_head_size)
                    )
                    _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                    if self.is_nested_encoder:
                        return _sdpa_out, None
                    return self.out_proj(_sdpa_out), None
                except Exception:
                    pass

            scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
            if is_causal:
                causal = (
                    torch.triu(
                        torch.full((L, L_kv), float("-inf"), device=device, dtype=torch.float32),
                        diagonal=1,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal

            valid = self._normalize_valid_mask(eff_mask, B, L_kv, device)
            if valid is not None:
                scores = scores.masked_fill(~valid[:, None, None, :], -1e4)

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            attn_weights = F.softmax(scores, dim=-1).to(dtype)
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            attn_weights = self.dropout(attn_weights)

            # FASA-FIX-10: unified head_mask handling
            if head_mask is not None:
                hm = head_mask.to(attn_weights.dtype)
                if hm.dim() == 1:
                    hm = hm.view(1, -1, 1, 1)
                elif hm.dim() == 2:
                    hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
                elif hm.dim() == 3:
                    hm = hm.unsqueeze(1)
                attn_weights = attn_weights * hm

            out = (
                torch.matmul(attn_weights, v)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.all_head_size)
            )
            final_attn_weights = attn_weights if output_attentions else None
        else:
            # -------- FASA pruning path --------
            valid = self._normalize_valid_mask(eff_mask, B, L_kv, device)
            k_p, v_p, p_mask, kept_idx = self._fasa_prune(k, v, valid)

            attn_scores = (
                torch.matmul(q.float(), k_p.float().transpose(-2, -1).contiguous()) * scale
            )

            if p_mask is not None:
                attn_scores = attn_scores.masked_fill(~p_mask[:, None, None, :], -1e4)

            attn_scores = torch.nan_to_num(attn_scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            attn_weights = F.softmax(attn_scores, dim=-1).to(dtype)
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            attn_weights = self.dropout(attn_weights)

            if head_mask is not None:
                hm = head_mask.to(attn_weights.dtype)
                if hm.dim() == 1:
                    hm = hm.view(1, -1, 1, 1)
                elif hm.dim() == 2:
                    hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
                elif hm.dim() == 3:
                    hm = hm.unsqueeze(1)
                attn_weights = attn_weights * hm

            out = (
                torch.matmul(attn_weights, v_p)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.all_head_size)
            )

            if output_attentions and kept_idx is not None:
                full_weights = torch.zeros(B, H, L, L_kv, device=device, dtype=dtype)
                kept_exp = kept_idx.unsqueeze(1).unsqueeze(2).expand(-1, H, L, -1)
                final_attn_weights = full_weights.scatter(3, kept_exp, attn_weights)
            elif output_attentions:
                final_attn_weights = attn_weights
            else:
                final_attn_weights = None

        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)

        return output, final_attn_weights

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if p is None:
                return 0
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("W_q", self.W_q),
            ("W_k", self.W_k),
            ("W_v", self.W_v),
            ("out_proj", self.out_proj),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["total"] = sum(out.values())
        return out

MixAttention-MoSKA (FASA): frequency/variance-aware sparse attention.

For non-causal self-attention and encoder-style models, we prune the K/V set down to the top fasa_keep_ratio tokens by variance importance, always preserving the [CLS] token at index 0.

For causal decoders, cross-attention, and incremental decoding (when past_key_value is present), we bypass pruning entirely and run exact dense attention via SDPA / manual softmax.

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.

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.W_k

prop k_proj

Expand source code

@property
def k_proj(self):
    return self.W_k

prop key

Expand source code

@property
def key(self):
    return self.W_k

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.W_q

prop q_proj

Expand source code

@property
def q_proj(self):
    return self.W_q

prop query

Expand source code

@property
def query(self):
    return self.W_q

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.W_v

prop v_proj

Expand source code

@property
def v_proj(self):
    return self.W_v

prop value

Expand source code

@property
def value(self):
    return self.W_v

Methods

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if p is None:
            return 0
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("W_q", self.W_q),
        ("W_k", self.W_k),
        ("W_v", self.W_v),
        ("out_proj", self.out_proj),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class MixtureOfExperts (dim: int,
ff_dim: int,
num_experts: int = 4,
top_k: int = 2,
**kwargs: Any)

Expand source code

class MixtureOfExperts(nn.Module):
    r"""Top-$k$ Token-Choice Mixture-of-Experts (MoE) Layer."""

    def __init__(
        self,
        dim: int,
        ff_dim: int,
        num_experts: int = 4,
        top_k: int = 2,
        **kwargs: Any,
    ) -> None:
        super().__init__()
        self.output_scale = nn.Parameter(torch.ones(1))
        self.dim: int = dim
        self.num_experts: int = num_experts
        self.top_k = top_k
        self.experts: nn.ModuleList = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(dim, ff_dim),
                    nn.GELU(),
                    nn.Linear(ff_dim, dim),
                )
                for _ in range(num_experts)
            ]
        )
        self.router: nn.Linear = nn.Linear(dim, num_experts)

    def forward(self, x: torch.Tensor, **kwargs: Any) -> torch.Tensor:
        B, L, D = x.shape
        scores = self.router(x)  # (B, L, E)
        k = min(self.top_k, scores.shape[-1])
        weights, indices = torch.topk(scores, k, dim=-1)
        weights = F.softmax(weights, dim=-1)
        out = torch.zeros_like(x)
        for e in range(self.num_experts):
            mask = (indices == e).any(dim=-1)  # (B, L)
            if mask.any():
                selected_weights = (weights * (indices == e).float()).sum(dim=-1, keepdim=True)
                out += (
                    mask.unsqueeze(-1) * selected_weights * self.experts[e](x * mask.unsqueeze(-1))
                )
        return out

Top-$k$ Token-Choice Mixture-of-Experts (MoE) Layer.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self,
x: torch.Tensor,
**kwargs: Any) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor, **kwargs: Any) -> torch.Tensor:
    B, L, D = x.shape
    scores = self.router(x)  # (B, L, E)
    k = min(self.top_k, scores.shape[-1])
    weights, indices = torch.topk(scores, k, dim=-1)
    weights = F.softmax(weights, dim=-1)
    out = torch.zeros_like(x)
    for e in range(self.num_experts):
        mask = (indices == e).any(dim=-1)  # (B, L)
        if mask.any():
            selected_weights = (weights * (indices == e).float()).sum(dim=-1, keepdim=True)
            out += (
                mask.unsqueeze(-1) * selected_weights * self.experts[e](x * mask.unsqueeze(-1))
            )
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class MoBAAttention (config: Any,
block_size: int = 64,
top_k: int = 4,
rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("moba", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("moba_attention", supports_causal=True, supports_gqa=True)
class MoBAAttention(BaseAttention):
    """
    MoBA: Mixture of Block Attention for Long-Context Modeling.

    For each query block, we compute block-level centroid scores, select the
    top-k most relevant key blocks, and run dense attention within the
    selected blocks (plus the causal-visible diagonal block).
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "sparse_attn"

    def __init__(
        self,
        config: Any,
        block_size: int = 64,
        top_k: int = 4,
        rank: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        # MB-FIX-03
        super().__init__(config, rank if rank is not None else 32)

        # MB-FIX-04: config None-safe
        if config is None:

            class _Cfg:
                pass

            config = _Cfg()
            config.hidden_size = kwargs.get("hidden_size", 768)
            config.num_attention_heads = kwargs.get("num_attention_heads", 12)
            config.head_dim = kwargs.get("head_dim", 64)
            config.model_type = kwargs.get("model_type", "")

        self.num_heads = int(
            getattr(
                config,
                "num_attention_heads",
                getattr(config, "num_heads", kwargs.get("num_attention_heads", 4)),
            )
        )
        self.hidden_size = int(getattr(config, "hidden_size", kwargs.get("hidden_size", 768)))
        cfg_hd = getattr(config, "head_dim", None)
        self.head_dim = (
            int(cfg_hd)
            if cfg_hd is not None
            else max(1, self.hidden_size // max(1, self.num_heads))
        )
        self.all_head_size = self.num_heads * self.head_dim

        cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = self.num_heads // self.num_kv_heads
        self.kv_head_size = self.num_kv_heads * self.head_dim

        self.block_size = int(block_size)
        self.top_k = int(top_k)
        self.lora_rank = int(lora_rank)
        self.lora_alpha = float(lora_alpha)

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(getattr(config, "model_type", "") or "").lower()
            self.is_decoder = bool(
                getattr(config, "is_decoder", False) or getattr(config, "causal", False)
            ) or any(
                _mt.startswith(p)
                for p in (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "falcon",
                    "bloom",
                    "opt",
                    "gemma",
                    "phi",
                    "olmo",
                    "mpt",
                    "stablelm",
                )
            )

        self.is_nested_encoder: bool = False

        # ---- MB-FIX-01: output scale (initialized to 1.0) ----
        self.output_scale = nn.Parameter(torch.ones(1))

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
        )
        self.dropout = nn.Dropout(dropout_p)

        # ---- projections ----
        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,
        )

        # MB-FIX-11: storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0
        self.triton_kwargs = {"block_size": self.block_size, "eps": 1e-6}

        pretrained_attn = kwargs.get("pretrained_attn")
        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    # ---- MB-FIX-07: property 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 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 c_proj(self):
        return self.out_proj

    @property
    def o_proj(self):
        return self.out_proj

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        # Track original storage
        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. UniversalAdapter
        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 as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # 2. GPT-2 fused c_attn
        if hasattr(src, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            c_attn = src.c_attn
            base_c = getattr(c_attn, "base_layer", c_attn)
            cw = getattr(base_c, "weight", getattr(c_attn, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                if cw.dim() == 2 and cw.shape[0] < cw.shape[1]:
                    cw = cw.t().contiguous()
                cb = getattr(base_c, "bias", getattr(c_attn, "bias", None))
                cb_data = cb.data.clone() if cb is not None else None
                split = cw.shape[0] // 3
                extracted["q"] = (
                    cw[:split],
                    cb_data[:split] if cb_data is not None else None,
                )
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb_data[split : 2 * split] if cb_data is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb_data[2 * split :] if cb_data is not None else None,
                )

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            cp = src.c_proj
            base_o = getattr(cp, "base_layer", cp)
            ow = getattr(base_o, "weight", getattr(cp, "weight", None))
            if ow is not None:
                ow = ow.data.clone()
                if ow.dim() == 2 and ow.shape[0] < ow.shape[1]:
                    ow = ow.t().contiguous()
                bo = getattr(base_o, "bias", getattr(cp, "bias", None))
                extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # 3. RoBERTa / BERT nested (MB-FIX-02: pass .data, not module)
        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") and extracted["q"][0] is None:
                q_l = getattr(s.query, "base_layer", s.query)
                extracted["q"] = (
                    q_l.weight.data,
                    q_l.bias.data if getattr(q_l, "bias", None) is not None else None,
                )
            if hasattr(s, "key") and extracted["k"][0] is None:
                k_l = getattr(s.key, "base_layer", s.key)
                extracted["k"] = (
                    k_l.weight.data,
                    k_l.bias.data if getattr(k_l, "bias", None) is not None else None,
                )
            if hasattr(s, "value") and extracted["v"][0] is None:
                v_l = getattr(s.value, "base_layer", s.value)
                extracted["v"] = (
                    v_l.weight.data,
                    v_l.bias.data if getattr(v_l, "bias", None) is not None else None,
                )

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            d = getattr(src.output.dense, "base_layer", src.output.dense)
            extracted["o"] = (
                d.weight.data,
                d.bias.data if getattr(d, "bias", None) is not None else None,
            )

        # 4. Modern decoders / DistilBERT
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(src, name):
                        m = getattr(src, name)
                        w = getattr(m, "weight", None)
                        if w is not None and hasattr(w, "data"):
                            b = getattr(m, "bias", None)
                            extracted[proj] = (
                                w.data,
                                b.data if b is not None else None,
                            )
                            break

        # MB-FIX-06: nested encoder detection after all extraction
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or cls_name in ("BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention")
            or cls_name.endswith("SelfAttention")
            and not has_out_proj
            or cls_name.endswith("Attention")
            and not has_out_proj
        ):
            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(
                "MoBA copy_weights: missing %s; proceeding with initialized weights.",
                missing,
            )

        # Frobenius norm checks
        for k_ in ("q", "k", "v"):
            w = extracted[k_][0]
            if w is not None:
                n = torch.linalg.norm(w.float()).item()
                if n < 1e-4:
                    logger.warning("MoBA: %s weight near-zero norm (%.2e)", k_, n)
        if extracted["o"][0] is not None:
            n = torch.linalg.norm(extracted["o"][0].float()).item()
            if n < 1e-4:
                logger.warning("MoBA: o weight near-zero norm (%.2e)", n)

        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

        # ---- Copy ----
        with torch.no_grad():
            for proj_key, target in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape

                    if proj_key in ("k", "v"):
                        if src_out == out_d:
                            target_w.data.copy_(w_oriented)
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b)
                        elif src_out > out_d and src_out % out_d == 0:
                            rep = src_out // out_d
                            target_w.data.copy_(w_oriented.view(out_d, rep, -1).mean(dim=1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.view(out_d, rep).mean(dim=1))
                        elif src_out < out_d and out_d % src_out == 0:
                            rep = out_d // src_out
                            target_w.data.copy_(w_oriented.repeat(rep, 1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.repeat(rep))
                        else:
                            m0, m1 = min(out_d, src_out), min(in_d, src_in)
                            target_w.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                            if b is not None and target.bias is not None:
                                mb = min(target.bias.shape[0], b.shape[0])
                                target.bias.data[:mb].copy_(b[:mb])
                    else:
                        m0, m1 = min(out_d, src_out), min(in_d, src_in)
                        target_w.data[:m0, :m1].copy_(w_oriented[:m0, :m1])
                        if b is not None and target.bias is not None:
                            mb = min(target.bias.shape[0], b.shape[0])
                            target.bias.data[:mb].copy_(b[:mb])
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                target.freeze_base_weights()

        # Track compressed storage
        try:
            for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # ---- is_decoder override ----
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)
        elif (
            cls_name in ("GPT2Attention", "GPT2Block")
            or hasattr(src, "c_attn")
            or "gpt" in cls_name.lower()
        ):
            self.is_decoder = True

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        L_kv = kv_input.shape[1]

        Q = self.q_proj(hidden_states).view(B, L, H, d).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, L_kv, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, L_kv, 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 and not is_cross:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        is_causal = bool(self.is_decoder and not is_cross and L > 1)
        scale = 1.0 / math.sqrt(d)

        eff_mask = encoder_attention_mask if is_cross else attention_mask

        # ---- Normalize mask to bool (B, L_kv) ----
        valid_mask: torch.Tensor | None = None
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    valid_mask = eff_mask
                elif eff_mask.dtype.is_floating_point:
                    valid_mask = eff_mask > -1e4
                else:
                    valid_mask = eff_mask > 0
            elif eff_mask.dim() == 3:
                m = eff_mask
                if m.dtype == torch.bool:
                    valid_mask = m.squeeze(1)
                elif m.dtype.is_floating_point:
                    valid_mask = (m > -1e4).squeeze(1)
                else:
                    valid_mask = (m > 0).squeeze(1)
            elif eff_mask.dim() == 4:
                m = eff_mask
                if m.dtype == torch.bool:
                    vv = m
                elif m.dtype.is_floating_point:
                    vv = m > -1e4
                else:
                    vv = m > 0
                valid_mask = vv.any(dim=-2).squeeze(1)

            if valid_mask is not None and valid_mask.shape[-1] != L_kv:
                valid_mask = None  # shape mismatch → ignore

        # ---- Fast path: short sequence or cross-attention ----
        if self.block_size >= L or is_cross:
            if valid_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 = self.is_decoder and not is_cross and L > 1
                    _sdpa_out = F.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, L, self.all_head_size)
                    )
                    _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                    if self.is_nested_encoder:
                        return _sdpa_out * self.output_scale, None
                    out = self.out_proj(_sdpa_out)
                    return out * self.output_scale, None
                except Exception:
                    pass

            # Dense manual
            scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
            if is_causal:
                causal_mask = (
                    torch.triu(
                        torch.full((L, L_kv), float("-inf"), device=device, dtype=torch.float32),
                        diagonal=1,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal_mask

            if valid_mask is not None:
                scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            attn_probs = F.softmax(scores, dim=-1).to(dtype)
            attn_probs = torch.nan_to_num(attn_probs, nan=0.0)
            attn_probs = self.dropout(attn_probs)

            # MB-FIX-09: head_mask shape
            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.shape[0], head_mask.shape[1], 1, 1)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(1)  # interpret as (B, H, 1)
                attn_probs = attn_probs * head_mask.to(attn_probs.dtype)

            context = torch.matmul(attn_probs, V)
            out = context.transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        else:
            # ---- Block-sparse path ----
            pad = (self.block_size - L % self.block_size) % self.block_size
            if pad > 0:
                Q_pad = F.pad(Q, (0, 0, 0, pad))
                K_pad = F.pad(K, (0, 0, 0, pad))
                V_pad = F.pad(V, (0, 0, 0, pad))
                if valid_mask is not None:
                    valid_mask_padded = F.pad(valid_mask, (0, pad), value=False)
                else:
                    valid_mask_padded = torch.ones(B, L + pad, device=device, dtype=torch.bool)
            else:
                Q_pad, K_pad, V_pad = Q, K, V
                if valid_mask is not None:
                    valid_mask_padded = valid_mask
                else:
                    valid_mask_padded = torch.ones(B, L, device=device, dtype=torch.bool)

            Lp = L + pad
            nb = Lp // self.block_size
            bs = self.block_size

            Qb = Q_pad.view(B, H, nb, bs, d)
            Kb = K_pad.view(B, H, nb, bs, d)
            Vb = V_pad.view(B, H, nb, bs, d)

            # ---- Block-level scores ----
            Q_mean = Qb.mean(dim=3)
            K_mean = Kb.mean(dim=3)
            block_scores = (
                torch.matmul(Q_mean.float(), K_mean.float().transpose(-2, -1).contiguous()) * scale
            )

            # MB-FIX-08: block-0 preservation
            # - For bidirectional encoders, block 0 (CLS) should always be a candidate.
            # - For causal decoders, causality is enforced per-position below, so we
            #   don't need to force block 0; a query at block i > 0 would never see
            #   block 0 through causal masking.
            if not is_causal:
                block_scores_for_topk = block_scores.clone()
                block_scores_for_topk[:, :, :, 0] = block_scores_for_topk[:, :, :, 0] + 1e4
            else:
                block_scores_for_topk = block_scores.clone()
                # For causal: a query in block i can only attend to blocks j ≤ i.
                # Mask blocks j > i to -inf.
                causal_block = torch.triu(
                    torch.full((nb, nb), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                block_scores_for_topk = block_scores_for_topk + causal_block.unsqueeze(0).unsqueeze(
                    0
                )

            top_k = min(self.top_k, nb)
            _, top_idx = block_scores_for_topk.topk(top_k, dim=-1)  # (B, H, nb, top_k)

            # MB-FIX-05: valid block mask reshaped correctly
            valid_b_exp = valid_mask_padded.view(B, 1, nb, bs).expand(-1, H, -1, -1)

            out_blocks = torch.zeros_like(Qb)
            for i in range(nb):
                idx = top_idx[:, :, i, :]  # (B, H, top_k)
                idx_exp = idx.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, -1, bs, d)
                K_sel = torch.gather(Kb, 2, idx_exp).view(B, H, top_k * bs, d)
                V_sel = torch.gather(Vb, 2, idx_exp).view(B, H, top_k * bs, d)

                idx_mask_exp = idx.unsqueeze(-1).expand(-1, -1, -1, bs)
                k_valid = torch.gather(valid_b_exp, 2, idx_mask_exp).reshape(B, H, 1, top_k * bs)

                Q_i = Qb[:, :, i]  # (B, H, bs, d)
                scores_i = (
                    torch.matmul(Q_i.float(), K_sel.float().transpose(-2, -1).contiguous()) * scale
                )
                scores_i = scores_i.masked_fill(~k_valid, -1e4)

                if is_causal:
                    q_pos = (i * bs + torch.arange(bs, device=device)).view(1, 1, bs, 1)
                    # absolute key positions for each of the top_k blocks
                    k_offsets = idx.unsqueeze(-1) * bs + torch.arange(bs, device=device).view(
                        1, 1, 1, bs
                    )
                    k_pos = k_offsets.reshape(B, H, 1, top_k * bs)
                    invalid = k_pos > q_pos
                    scores_i = scores_i.masked_fill(invalid, -1e4)

                scores_i = torch.nan_to_num(scores_i, nan=-1e4, posinf=1e4, neginf=-1e4)
                attn_i = F.softmax(scores_i, dim=-1).to(dtype)
                attn_i = torch.nan_to_num(attn_i, nan=0.0)
                attn_i = self.dropout(attn_i)
                out_blocks[:, :, i] = torch.matmul(attn_i, V_sel)

            out_pad = out_blocks.view(B, H, Lp, d)[:, :, :L, :]
            out = out_pad.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)

        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            final_out = out
        else:
            final_out = self.out_proj(out)

        # MB-FIX-01: apply output_scale
        final_out = final_out * self.output_scale

        return final_out, None

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        def _sz(p):
            if p is None:
                return 0
            if isinstance(p, (nn.Parameter, torch.Tensor)):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("q_proj", self.q_proj),
            ("k_proj", self.k_proj),
            ("v_proj", self.v_proj),
            ("out_proj", self.out_proj),
        ):
            out[name] = sum(_sz(p) for p in m.parameters())
        out["output_scale"] = _sz(self.output_scale)
        out["total"] = sum(out.values())
        return out

MoBA: Mixture of Block Attention for Long-Context Modeling.

For each query block, we compute block-level centroid scores, select the top-k most relevant key blocks, and run dense attention within the selected blocks (plus the causal-visible diagonal block).

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.

var triton_kernel : str

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 storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    def _sz(p):
        if p is None:
            return 0
        if isinstance(p, (nn.Parameter, torch.Tensor)):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("q_proj", self.q_proj),
        ("k_proj", self.k_proj),
        ("v_proj", self.v_proj),
        ("out_proj", self.out_proj),
    ):
        out[name] = sum(_sz(p) for p in m.parameters())
    out["output_scale"] = _sz(self.output_scale)
    out["total"] = sum(out.values())
    return out

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class MoEDeltaKVPagedAttention (config: Any | None = None,
page_size: int = 16,
k_experts: int = 4,
svd_rank: int = 32,
q_bits: int = 4,
max_delta_error: float = 0.35,
lora_rank: int = 16,
lora_alpha: float = 32.0,
enable_moe: bool = True,
enable_deltakv: bool = True,
is_decoder: bool | None = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("moe_deltakv_paged", supports_causal=True, supports_gqa=True)
class MoEDeltaKVPagedAttention(nn.Module):
    """
    MoE-DeltaKV-Paged Hybrid Attention.

    Two execution paths:
      * Full dense attention when `num_pages <= 1` (page_size >= seq_len)
        OR when `enable_moe=False`.
      * Paged MoE attention otherwise.
    """

    def __init__(
        self,
        config: Any | None = None,
        page_size: int = 16,
        k_experts: int = 4,
        svd_rank: int = 32,
        q_bits: int = 4,
        max_delta_error: float = 0.35,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        enable_moe: bool = True,
        enable_deltakv: bool = True,
        is_decoder: bool | None = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__()

        if config is None:

            class _Cfg:
                pass

            config = _Cfg()
            config.hidden_size = kwargs.get("hidden_size", 768)
            config.num_attention_heads = kwargs.get("num_attention_heads", 12)
            config.head_dim = kwargs.get("head_dim", 64)
            config.model_type = kwargs.get("model_type", "")
            config.is_decoder = kwargs.get("is_decoder", False)

        # ---- dim inference (MDP-FIX-14) ----
        self.hidden_size = int(
            getattr(
                config,
                "hidden_size",
                getattr(config, "d_model", kwargs.get("hidden_size", 768)),
            )
        )
        self.num_heads = int(
            getattr(
                config,
                "num_attention_heads",
                getattr(config, "num_heads", kwargs.get("num_attention_heads", 12)),
            )
        )
        self.head_dim = int(
            getattr(
                config,
                "head_dim",
                kwargs.get("head_dim", self.hidden_size // max(1, self.num_heads)),
            )
        )

        cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = max(1, self.num_heads // self.num_kv_heads)

        self.all_head_size = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim

        # ---- paging config (MDP-FIX-01, MDP-FIX-12) ----
        self.page_size = max(4, int(page_size))
        self.k_experts = max(1, int(k_experts))
        self.svd_rank = max(4, int(svd_rank))
        self.q_bits = int(q_bits)
        if self.q_bits not in (2, 4, 8):
            self.q_bits = 4
        self.max_delta_error = float(max_delta_error)

        # ---- flags ----
        self.enable_moe = bool(enable_moe)
        self.enable_deltakv = bool(enable_deltakv)
        self.lora_rank = max(0, int(lora_rank))
        self.lora_alpha = float(lora_alpha)

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_causal = bool(is_decoder)
        else:
            _mt = str(getattr(config, "model_type", "") or "").lower()
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            self.is_causal = bool(
                getattr(config, "is_causal", False)
                or getattr(config, "is_decoder", False)
                or kwargs.get("is_causal", False)
                or kwargs.get("is_decoder", False)
                or any(_mt.startswith(p) for p in _dec_pfx)
            )
        self.is_decoder = self.is_causal
        self.is_nested_encoder = False

        # ---- projections ----
        self.query = LoRALinear(
            self.hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.key = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.value = 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,
        )

        # ---- page router (identity init so it preserves original attention) ----
        self.W_rg = nn.Linear(self.head_dim, self.head_dim, bias=False)
        with torch.no_grad():
            nn.init.eye_(self.W_rg.weight)

        # ---- auxiliary load-balancing loss as a buffer (MDP-FIX-03) ----
        self.register_buffer("aux_loss", torch.zeros(1), persistent=False)

        # ---- delta-KV reconstruction telemetry ----
        self.register_buffer("delta_kv_error_k", torch.zeros(1), persistent=False)
        self.register_buffer("delta_kv_error_v", torch.zeros(1), persistent=False)

        dropout_p = float(
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                ),
            )
        )
        self.dropout = nn.Dropout(dropout_p)

        # ---- storage accounting (MDP-FIX-07) ----
        self._orig_storage_bytes: int = 0
        self._comp_storage_bytes: int = 0
        self._source_id: int | None = None

        src = pretrained_attn or kwargs.get("old_attn")
        if src is not None:
            self.copy_weights(src)

    # ------------------------------------------------------------------
    @property
    def q_proj(self):
        return self.query

    @property
    def k_proj(self):
        return self.key

    @property
    def v_proj(self):
        return self.value

    @property
    def q_lin(self):
        return self.query

    @property
    def k_lin(self):
        return self.key

    @property
    def v_lin(self):
        return self.value

    @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 storage_bytes(self) -> dict[str, int]:
        return {"orig": self._orig_storage_bytes, "comp": self._comp_storage_bytes}

    # ------------------------------------------------------------------
    # copy_weights (MDP-FIX-05, MDP-FIX-09)
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        if src is None:
            return

        from neuralops.core.adapters import UniversalAdapter

        # Idempotency
        if self._source_id is not None and id(src) == self._source_id:
            return

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # --- 1. UniversalAdapter ---
        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 ---
        if hasattr(src, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            try:
                mod_c = src.c_attn
                base_c = getattr(mod_c, "base_layer", mod_c)
                weight = getattr(base_c, "weight", getattr(mod_c, "weight", None))
                if weight is not None:
                    weight = weight.data.clone()
                    bias = getattr(base_c, "bias", getattr(mod_c, "bias", None))
                    bias_data = bias.data.clone() if bias is not None else None
                    cls_c = mod_c.__class__.__name__
                    base_cls = base_c.__class__.__name__
                    if (
                        hasattr(mod_c, "nf")
                        or hasattr(base_c, "nf")
                        or "Conv1D" in cls_c
                        or "Conv1D" in base_cls
                        or (weight.dim() == 2 and weight.shape[0] < weight.shape[1])
                    ):
                        weight = weight.t().contiguous()
                    split = weight.shape[0] // 3
                    extracted["q"] = (
                        weight[:split],
                        bias_data[:split] if bias_data is not None else None,
                    )
                    extracted["k"] = (
                        weight[split : 2 * split],
                        bias_data[split : 2 * split] if bias_data is not None else None,
                    )
                    extracted["v"] = (
                        weight[2 * split :],
                        bias_data[2 * split :] if bias_data is not None else None,
                    )
            except Exception as e:
                logger.debug("GPT-2 c_attn extraction failed: %s", e)

        if hasattr(src, "c_proj") and extracted["o"][0] is None:
            try:
                mod_o = src.c_proj
                base_o = getattr(mod_o, "base_layer", mod_o)
                w_o = getattr(base_o, "weight", getattr(mod_o, "weight", None))
                if w_o is not None:
                    w_o = w_o.data.clone()
                    cls_o = mod_o.__class__.__name__
                    base_cls_o = base_o.__class__.__name__
                    if (
                        hasattr(mod_o, "nf")
                        or hasattr(base_o, "nf")
                        or "Conv1D" in cls_o
                        or "Conv1D" in base_cls_o
                        or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
                    ):
                        w_o = w_o.t().contiguous()
                    b_o = getattr(base_o, "bias", getattr(mod_o, "bias", None))
                    extracted["o"] = (
                        w_o,
                        b_o.data.clone() if b_o is not None else None,
                    )
            except Exception as e:
                logger.debug("GPT-2 c_proj extraction failed: %s", e)

        # --- 3. BERT / RoBERTa nested ---
        if hasattr(src, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
            try:
                s = src.self
                if hasattr(s, "query") and extracted["q"][0] is None:
                    extracted["q"] = (
                        s.query.weight.data.clone(),
                        s.query.bias.data.clone() if s.query.bias is not None else None,
                    )
                if hasattr(s, "key") and extracted["k"][0] is None:
                    extracted["k"] = (
                        s.key.weight.data.clone(),
                        s.key.bias.data.clone() if s.key.bias is not None else None,
                    )
                if hasattr(s, "value") and extracted["v"][0] is None:
                    extracted["v"] = (
                        s.value.weight.data.clone(),
                        s.value.bias.data.clone() if s.value.bias is not None else None,
                    )
            except Exception as e:
                logger.debug("Nested self extraction failed: %s", e)

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            try:
                dense = src.output.dense
                extracted["o"] = (
                    dense.weight.data.clone(),
                    dense.bias.data.clone() if dense.bias is not None else None,
                )
            except Exception as e:
                logger.debug("output.dense extraction failed: %s", e)

        # --- 4. Flat modules ---
        named_map = {
            "q": ("q_proj", "q_lin", "query"),
            "k": ("k_proj", "k_lin", "key"),
            "v": ("v_proj", "v_lin", "value"),
            "o": ("o_proj", "out_proj", "out_lin", "dense"),
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(src, name):
                        try:
                            layer = getattr(src, name)
                            if hasattr(layer, "weight"):
                                extracted[proj] = (
                                    layer.weight.data.clone(),
                                    (
                                        layer.bias.data.clone()
                                        if getattr(layer, "bias", None) is not None
                                        else None
                                    ),
                                )
                                break
                        except Exception:
                            continue

        # --- 5. Nested encoder detection (MDP-FIX-05) ---
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, a)
            for a in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or cls_name
            in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
            )
        ):
            is_nested = True
        elif cls_name in ("DistilBertAttention", "DistilBertSABlock"):
            # DistilBERT nests attention inside a container but has an out_lin
            is_nested = True
        elif (
            cls_name.endswith("SelfAttention")
            and not has_out_proj
            or cls_name.endswith("Attention")
            and not has_out_proj
        ):
            is_nested = True
        self.is_nested_encoder = bool(is_nested)

        # --- 6. Missing check ---
        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("MoEDeltaKVPaged: missing %s; aborting copy.", missing)
            return

        # --- 7. Storage tracking ---
        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # --- 8. Copy into projections ---
        with torch.no_grad():
            # Q, K, V
            for proj_key, target in (
                ("q", self.query),
                ("k", self.key),
                ("v", self.value),
            ):
                w, b = extracted[proj_key]
                if w is None:
                    continue
                in_dim = self.hidden_size
                out_dim = target.weight.shape[0]

                if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim:
                    w = w.t().contiguous()

                # GQA expansion (repeat)
                if w.shape[0] < out_dim and out_dim % w.shape[0] == 0:
                    rep = out_dim // w.shape[0]
                    w = w.repeat(rep, 1)
                    if b is not None:
                        b = b.repeat(rep)
                elif w.shape[0] > out_dim and w.shape[0] % out_dim == 0:
                    rep = w.shape[0] // out_dim
                    w = w.view(out_dim, rep, -1).mean(dim=1)
                    if b is not None:
                        b = b.view(out_dim, rep).mean(dim=1)

                target.weight.data.zero_()
                out_rows = min(out_dim, w.shape[0])
                in_cols = min(in_dim, w.shape[1])
                target.weight.data[:out_rows, :in_cols].copy_(
                    w[:out_rows, :in_cols].to(target.weight.device, target.weight.dtype)
                )
                if b is not None and target.bias is not None:
                    mb = min(target.bias.numel(), b.numel())
                    target.bias.data.zero_()
                    target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))

            # Output projection
            w_o, b_o = extracted["o"]
            if w_o is not None and not self.is_nested_encoder:
                in_dim = self.all_head_size
                out_dim = self.hidden_size
                if w_o.dim() == 2 and w_o.shape[1] != in_dim and w_o.shape[0] == in_dim:
                    w_o = w_o.t().contiguous()
                self.out_proj.weight.data.zero_()
                out_rows = min(out_dim, w_o.shape[0])
                in_cols = min(in_dim, w_o.shape[1])
                self.out_proj.weight.data[:out_rows, :in_cols].copy_(
                    w_o[:out_rows, :in_cols].to(
                        self.out_proj.weight.device, self.out_proj.weight.dtype
                    )
                )
                if b_o is not None and self.out_proj.bias is not None:
                    mb = min(self.out_proj.bias.numel(), b_o.numel())
                    self.out_proj.bias.data.zero_()
                    self.out_proj.bias.data[:mb].copy_(
                        b_o[:mb].to(self.out_proj.bias.device, self.out_proj.bias.dtype)
                    )
            elif self.is_nested_encoder:
                # Nested encoder: out_proj is bypassed in forward.
                pass

            # Freeze base weights
            for m in (self.query, self.key, self.value, self.out_proj):
                m.freeze_base_weights()

        # --- 9. Storage accounting ---
        try:
            for m in (self.query, self.key, self.value, self.out_proj):
                for p in m.parameters():
                    self._comp_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        # --- 10. is_decoder override ---
        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)
            self.is_causal = self.is_decoder

        self._source_id = id(src)

    # ------------------------------------------------------------------
    # Paged MoE block attention (executed when num_pages >= 2)
    # ------------------------------------------------------------------
    def _paged_moe_attention(
        self,
        q: torch.Tensor,
        k_hat: torch.Tensor,
        v_hat: torch.Tensor,
        valid_mask: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, H, N, d = q.shape
        page_size = self.page_size
        num_pages = (N + page_size - 1) // page_size
        pad_len = num_pages * page_size - N

        if pad_len > 0:
            q_pad = F.pad(q, (0, 0, 0, pad_len))
            k_pad = F.pad(k_hat, (0, 0, 0, pad_len))
            v_pad = F.pad(v_hat, (0, 0, 0, pad_len))
            valid_pad = F.pad(valid_mask, (0, pad_len), value=False)
        else:
            q_pad, k_pad, v_pad, valid_pad = q, k_hat, v_hat, valid_mask

        padded_N = num_pages * page_size

        # ---- Blockify ----
        q_blocks = q_pad.view(B, H, num_pages, page_size, d)
        k_blocks = k_pad.view(B, H, num_pages, page_size, d)
        v_blocks = v_pad.view(B, H, num_pages, page_size, d)

        valid_blocks = valid_pad.view(B, 1, num_pages, page_size).float()

        # ---- Page-level pooling ----
        valid_sum = valid_blocks.sum(dim=-1, keepdim=True).clamp_min(1.0)
        q_bar = (q_blocks * valid_blocks.unsqueeze(-1)).sum(dim=3) / valid_sum
        k_bar = (k_blocks * valid_blocks.unsqueeze(-1)).sum(dim=3) / valid_sum

        # ---- Routing ----
        router_logits = torch.matmul(
            self.W_rg(q_bar), k_bar.transpose(-1, -2).contiguous()
        ) / math.sqrt(d)

        # Page validity: at least one valid token in the page
        page_valid = valid_pad.view(B, 1, num_pages, page_size).any(dim=-1)  # (B,1,num_pages)
        router_logits = router_logits.masked_fill(~page_valid.unsqueeze(2), float("-inf"))

        # Causal masking of pages
        if self.is_causal:
            causal_pages = torch.tril(
                torch.ones(num_pages, num_pages, device=q.device, dtype=torch.bool)
            )
            router_logits = router_logits.masked_fill(~causal_pages, float("-inf"))

        router_logits = torch.nan_to_num(router_logits, nan=-1e4, posinf=1e4, neginf=-1e4)

        # Ensure page 0 always selected (attention sink / CLS protection)
        router_logits_for_topk = router_logits.clone()
        if self.is_causal:
            # Only bias page 0 for queries whose page index >= 0 (always true)
            router_logits_for_topk[:, :, :, 0] = router_logits_for_topk[:, :, :, 0] + 1e4
        else:
            router_logits_for_topk[:, :, :, 0] = router_logits_for_topk[:, :, :, 0] + 1e4

        k_experts = min(self.k_experts, num_pages)
        _, topk_indices = torch.topk(router_logits_for_topk, k_experts, dim=-1)
        topk_logits = torch.gather(router_logits, -1, topk_indices)
        g_j = F.softmax(topk_logits, dim=-1)  # (B,H,num_pages,k)

        # ---- Auxiliary loss (load balancing) ----
        router_probs = F.softmax(router_logits, dim=-1)  # (B,H,num_pages,num_pages)
        expert_mask = F.one_hot(
            topk_indices, num_classes=num_pages
        ).float()  # (B,H,num_pages,k,num_pages)
        density = expert_mask.sum(dim=-2).mean(dim=(0, 1, 2))  # (num_pages,)
        prob_mean = router_probs.mean(dim=(0, 1, 2))  # (num_pages,)
        with torch.no_grad():
            self.aux_loss.copy_(num_pages * torch.sum(density * prob_mean).detach())

        # ---- Gather selected pages ----
        expanded_idx = topk_indices.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, -1, -1, page_size, d)
        k_blocks_exp = k_blocks.unsqueeze(2).expand(-1, -1, num_pages, -1, -1, -1)
        v_blocks_exp = v_blocks.unsqueeze(2).expand(-1, -1, num_pages, -1, -1, -1)
        k_gathered = torch.gather(
            k_blocks_exp, dim=3, index=expanded_idx
        )  # (B,H,num_pages,k,page_size,d)
        v_gathered = torch.gather(v_blocks_exp, dim=3, index=expanded_idx)

        # ---- Gather validity mask ----
        valid_blocks_exp = valid_blocks.unsqueeze(2).expand(-1, -1, num_pages, -1, -1)
        exp_valid_idx = topk_indices.unsqueeze(-1).expand(-1, -1, -1, -1, page_size)
        valid_blocks_exp_h = valid_blocks_exp.expand(-1, H, -1, -1, -1)
        k_valid = torch.gather(
            valid_blocks_exp_h, dim=3, index=exp_valid_idx
        )  # (B,H,num_pages,k,page_size)

        # ---- Per-expert attention ----
        q_exp = q_blocks.unsqueeze(3)  # (B,H,num_pages,1,page_size,d)
        attn_scores = torch.matmul(q_exp, k_gathered.transpose(-1, -2).contiguous()) / math.sqrt(d)
        attn_scores = attn_scores.masked_fill(~k_valid.unsqueeze(-2), -1e4)
        attn_scores = torch.nan_to_num(attn_scores, nan=-1e4, posinf=1e4, neginf=-1e4)

        attn_probs = F.softmax(attn_scores, dim=-1)
        attn_probs = self.dropout(attn_probs)

        context_per_expert = torch.matmul(attn_probs, v_gathered)  # (B,H,num_pages,k,page_size,d)

        # ---- Mix experts with gate weights ----
        g_j_exp = g_j.unsqueeze(-1).unsqueeze(-1)  # (B,H,num_pages,k,1,1)
        context = (context_per_expert * g_j_exp).sum(dim=3)  # (B,H,num_pages,page_size,d)

        # ---- Reshape to (B,H,padded_N,d) ----
        context = context.reshape(B, H, padded_N, d)

        return context, attn_probs

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:

        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, N, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        H_kv = self.num_kv_heads
        device = hidden_states.device
        dtype = hidden_states.dtype

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        kv_len = kv_input.shape[1]

        # ---- Projections ----
        q = self.query(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous()
        k = self.key(kv_input).view(B, kv_len, H_kv, d).transpose(1, 2).contiguous()
        v = self.value(kv_input).view(B, kv_len, H_kv, d).transpose(1, 2).contiguous()

        # ---- RoPE ----
        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
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # ---- GQA ----
        if H_kv != H:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        # ---- Delta-KV compression (MDP-FIX-02) ----
        if self.enable_deltakv and kv_len > 1:
            k_hat, v_hat, err_k, err_v = simulate_paged_delta_kv(
                k,
                v,
                rank=self.svd_rank,
                bits=self.q_bits,
                max_delta_error=self.max_delta_error,
            )
            with torch.no_grad():
                self.delta_kv_error_k.copy_(torch.tensor(err_k, device=device))
                self.delta_kv_error_v.copy_(torch.tensor(err_v, device=device))
        else:
            k_hat, v_hat = k, v

        # ---- Mask normalization ----
        valid_mask: torch.Tensor | None = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                if attention_mask.dtype == torch.bool:
                    valid_mask = attention_mask
                elif attention_mask.dtype.is_floating_point:
                    valid_mask = attention_mask > -1e4
                else:
                    valid_mask = attention_mask > 0
            elif attention_mask.dim() == 3:
                m = attention_mask
                valid_mask = (m > -1e4).squeeze(1) if m.dtype.is_floating_point else m.squeeze(1)
            elif attention_mask.dim() == 4:
                m = attention_mask
                v_bool = (
                    m
                    if m.dtype == torch.bool
                    else (m > -1e4 if m.dtype.is_floating_point else m > 0)
                )
                valid_mask = v_bool.any(dim=-2).squeeze(1)
            if valid_mask is not None and valid_mask.shape[-1] != kv_len:
                valid_mask = None

        if valid_mask is None:
            valid_mask = torch.ones(B, kv_len, device=device, dtype=torch.bool)

        # ---- Decide execution path ----
        num_pages = (kv_len + self.page_size - 1) // self.page_size
        use_paged = self.enable_moe and num_pages > 1 and not is_cross

        if use_paged:
            context, attn_probs = self._paged_moe_attention(q, k_hat, v_hat, valid_mask)
        else:
            # ---- Dense path ----
            scale = 1.0 / math.sqrt(d)
            scores = torch.matmul(q.float(), k_hat.float().transpose(-2, -1).contiguous()) * scale

            if self.is_causal and not is_cross and N > 1:
                diag = 1 + max(0, kv_len - N)
                causal = (
                    torch.triu(
                        torch.full(
                            (N, kv_len),
                            float("-inf"),
                            device=device,
                            dtype=torch.float32,
                        ),
                        diagonal=diag,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal

            scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)
            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)

            attn_probs = F.softmax(scores, dim=-1).to(dtype)
            attn_probs = self.dropout(attn_probs)
            context = torch.matmul(attn_probs, v_hat)

            with torch.no_grad():
                self.aux_loss.zero_()

        # ---- head_mask ----
        if head_mask is not None:
            hm = head_mask.to(context.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            context = context * hm

        # ---- Reshape to (B, N, all_head_size) ----
        context = context[:, :, :N, :] if context.shape[2] != N else context
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # ---- Output projection ----
        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_proj(context)

        # MDP-FIX-13: always return a 2-tuple with attn = tensor or None
        attn_ret = attn_probs if output_attentions else None
        return out, attn_ret

MoE-DeltaKV-Paged Hybrid Attention.

Two execution paths: * Full dense attention when num_pages <= 1 (page_size >= seq_len) OR when enable_moe=False. * Paged MoE attention otherwise.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

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.key

prop k_proj

Expand source code

@property
def k_proj(self):
    return self.key

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.query

prop q_proj

Expand source code

@property
def q_proj(self):
    return self.query

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.value

prop v_proj

Expand source code

@property
def v_proj(self):
    return self.value

Methods

def copy_weights(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    if src is None:
        return

    from neuralops.core.adapters import UniversalAdapter

    # Idempotency
    if self._source_id is not None and id(src) == self._source_id:
        return

    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

    # --- 1. UniversalAdapter ---
    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 ---
    if hasattr(src, "c_attn") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        try:
            mod_c = src.c_attn
            base_c = getattr(mod_c, "base_layer", mod_c)
            weight = getattr(base_c, "weight", getattr(mod_c, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(mod_c, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                cls_c = mod_c.__class__.__name__
                base_cls = base_c.__class__.__name__
                if (
                    hasattr(mod_c, "nf")
                    or hasattr(base_c, "nf")
                    or "Conv1D" in cls_c
                    or "Conv1D" in base_cls
                    or (weight.dim() == 2 and weight.shape[0] < weight.shape[1])
                ):
                    weight = weight.t().contiguous()
                split = weight.shape[0] // 3
                extracted["q"] = (
                    weight[:split],
                    bias_data[:split] if bias_data is not None else None,
                )
                extracted["k"] = (
                    weight[split : 2 * split],
                    bias_data[split : 2 * split] if bias_data is not None else None,
                )
                extracted["v"] = (
                    weight[2 * split :],
                    bias_data[2 * split :] if bias_data is not None else None,
                )
        except Exception as e:
            logger.debug("GPT-2 c_attn extraction failed: %s", e)

    if hasattr(src, "c_proj") and extracted["o"][0] is None:
        try:
            mod_o = src.c_proj
            base_o = getattr(mod_o, "base_layer", mod_o)
            w_o = getattr(base_o, "weight", getattr(mod_o, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                cls_o = mod_o.__class__.__name__
                base_cls_o = base_o.__class__.__name__
                if (
                    hasattr(mod_o, "nf")
                    or hasattr(base_o, "nf")
                    or "Conv1D" in cls_o
                    or "Conv1D" in base_cls_o
                    or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                b_o = getattr(base_o, "bias", getattr(mod_o, "bias", None))
                extracted["o"] = (
                    w_o,
                    b_o.data.clone() if b_o is not None else None,
                )
        except Exception as e:
            logger.debug("GPT-2 c_proj extraction failed: %s", e)

    # --- 3. BERT / RoBERTa nested ---
    if hasattr(src, "self") and any(extracted[k][0] is None for k in ("q", "k", "v")):
        try:
            s = src.self
            if hasattr(s, "query") and extracted["q"][0] is None:
                extracted["q"] = (
                    s.query.weight.data.clone(),
                    s.query.bias.data.clone() if s.query.bias is not None else None,
                )
            if hasattr(s, "key") and extracted["k"][0] is None:
                extracted["k"] = (
                    s.key.weight.data.clone(),
                    s.key.bias.data.clone() if s.key.bias is not None else None,
                )
            if hasattr(s, "value") and extracted["v"][0] is None:
                extracted["v"] = (
                    s.value.weight.data.clone(),
                    s.value.bias.data.clone() if s.value.bias is not None else None,
                )
        except Exception as e:
            logger.debug("Nested self extraction failed: %s", e)

    if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
        try:
            dense = src.output.dense
            extracted["o"] = (
                dense.weight.data.clone(),
                dense.bias.data.clone() if dense.bias is not None else None,
            )
        except Exception as e:
            logger.debug("output.dense extraction failed: %s", e)

    # --- 4. Flat modules ---
    named_map = {
        "q": ("q_proj", "q_lin", "query"),
        "k": ("k_proj", "k_lin", "key"),
        "v": ("v_proj", "v_lin", "value"),
        "o": ("o_proj", "out_proj", "out_lin", "dense"),
    }
    for proj, names in named_map.items():
        if extracted[proj][0] is None:
            for name in names:
                if hasattr(src, name):
                    try:
                        layer = getattr(src, name)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data.clone(),
                                (
                                    layer.bias.data.clone()
                                    if getattr(layer, "bias", None) is not None
                                    else None
                                ),
                            )
                            break
                    except Exception:
                        continue

    # --- 5. Nested encoder detection (MDP-FIX-05) ---
    cls_name = src.__class__.__name__
    has_out_proj = any(
        hasattr(src, a)
        for a in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    is_nested = False
    if (
        hasattr(src, "self")
        and hasattr(src.self, "query")
        or cls_name
        in (
            "BertSelfAttention",
            "RobertaSelfAttention",
            "XLMRobertaSelfAttention",
            "ElectraSelfAttention",
        )
    ):
        is_nested = True
    elif cls_name in ("DistilBertAttention", "DistilBertSABlock"):
        # DistilBERT nests attention inside a container but has an out_lin
        is_nested = True
    elif (
        cls_name.endswith("SelfAttention")
        and not has_out_proj
        or cls_name.endswith("Attention")
        and not has_out_proj
    ):
        is_nested = True
    self.is_nested_encoder = bool(is_nested)

    # --- 6. Missing check ---
    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("MoEDeltaKVPaged: missing %s; aborting copy.", missing)
        return

    # --- 7. Storage tracking ---
    try:
        for p in src.parameters():
            self._orig_storage_bytes += p.numel() * p.element_size()
    except Exception:
        pass

    # --- 8. Copy into projections ---
    with torch.no_grad():
        # Q, K, V
        for proj_key, target in (
            ("q", self.query),
            ("k", self.key),
            ("v", self.value),
        ):
            w, b = extracted[proj_key]
            if w is None:
                continue
            in_dim = self.hidden_size
            out_dim = target.weight.shape[0]

            if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim:
                w = w.t().contiguous()

            # GQA expansion (repeat)
            if w.shape[0] < out_dim and out_dim % w.shape[0] == 0:
                rep = out_dim // w.shape[0]
                w = w.repeat(rep, 1)
                if b is not None:
                    b = b.repeat(rep)
            elif w.shape[0] > out_dim and w.shape[0] % out_dim == 0:
                rep = w.shape[0] // out_dim
                w = w.view(out_dim, rep, -1).mean(dim=1)
                if b is not None:
                    b = b.view(out_dim, rep).mean(dim=1)

            target.weight.data.zero_()
            out_rows = min(out_dim, w.shape[0])
            in_cols = min(in_dim, w.shape[1])
            target.weight.data[:out_rows, :in_cols].copy_(
                w[:out_rows, :in_cols].to(target.weight.device, target.weight.dtype)
            )
            if b is not None and target.bias is not None:
                mb = min(target.bias.numel(), b.numel())
                target.bias.data.zero_()
                target.bias.data[:mb].copy_(b[:mb].to(target.bias.device, target.bias.dtype))

        # Output projection
        w_o, b_o = extracted["o"]
        if w_o is not None and not self.is_nested_encoder:
            in_dim = self.all_head_size
            out_dim = self.hidden_size
            if w_o.dim() == 2 and w_o.shape[1] != in_dim and w_o.shape[0] == in_dim:
                w_o = w_o.t().contiguous()
            self.out_proj.weight.data.zero_()
            out_rows = min(out_dim, w_o.shape[0])
            in_cols = min(in_dim, w_o.shape[1])
            self.out_proj.weight.data[:out_rows, :in_cols].copy_(
                w_o[:out_rows, :in_cols].to(
                    self.out_proj.weight.device, self.out_proj.weight.dtype
                )
            )
            if b_o is not None and self.out_proj.bias is not None:
                mb = min(self.out_proj.bias.numel(), b_o.numel())
                self.out_proj.bias.data.zero_()
                self.out_proj.bias.data[:mb].copy_(
                    b_o[:mb].to(self.out_proj.bias.device, self.out_proj.bias.dtype)
                )
        elif self.is_nested_encoder:
            # Nested encoder: out_proj is bypassed in forward.
            pass

        # Freeze base weights
        for m in (self.query, self.key, self.value, self.out_proj):
            m.freeze_base_weights()

    # --- 9. Storage accounting ---
    try:
        for m in (self.query, self.key, self.value, self.out_proj):
            for p in m.parameters():
                self._comp_storage_bytes += p.numel() * p.element_size()
    except Exception:
        pass

    # --- 10. is_decoder override ---
    if hasattr(src, "is_decoder"):
        self.is_decoder = bool(src.is_decoder)
        self.is_causal = self.is_decoder

    self._source_id = id(src)

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:

    if not hidden_states.is_contiguous():
        hidden_states = hidden_states.contiguous()

    B, N, D = hidden_states.shape
    H, d = self.num_heads, self.head_dim
    H_kv = self.num_kv_heads
    device = hidden_states.device
    dtype = hidden_states.dtype

    is_cross = encoder_hidden_states is not None
    kv_input = encoder_hidden_states if is_cross else hidden_states
    kv_len = kv_input.shape[1]

    # ---- Projections ----
    q = self.query(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous()
    k = self.key(kv_input).view(B, kv_len, H_kv, d).transpose(1, 2).contiguous()
    v = self.value(kv_input).view(B, kv_len, H_kv, d).transpose(1, 2).contiguous()

    # ---- RoPE ----
    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
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    # ---- GQA ----
    if H_kv != H:
        k = repeat_kv(k, self.kv_repeat)
        v = repeat_kv(v, self.kv_repeat)

    # ---- Delta-KV compression (MDP-FIX-02) ----
    if self.enable_deltakv and kv_len > 1:
        k_hat, v_hat, err_k, err_v = simulate_paged_delta_kv(
            k,
            v,
            rank=self.svd_rank,
            bits=self.q_bits,
            max_delta_error=self.max_delta_error,
        )
        with torch.no_grad():
            self.delta_kv_error_k.copy_(torch.tensor(err_k, device=device))
            self.delta_kv_error_v.copy_(torch.tensor(err_v, device=device))
    else:
        k_hat, v_hat = k, v

    # ---- Mask normalization ----
    valid_mask: torch.Tensor | None = None
    if attention_mask is not None:
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                valid_mask = attention_mask
            elif attention_mask.dtype.is_floating_point:
                valid_mask = attention_mask > -1e4
            else:
                valid_mask = attention_mask > 0
        elif attention_mask.dim() == 3:
            m = attention_mask
            valid_mask = (m > -1e4).squeeze(1) if m.dtype.is_floating_point else m.squeeze(1)
        elif attention_mask.dim() == 4:
            m = attention_mask
            v_bool = (
                m
                if m.dtype == torch.bool
                else (m > -1e4 if m.dtype.is_floating_point else m > 0)
            )
            valid_mask = v_bool.any(dim=-2).squeeze(1)
        if valid_mask is not None and valid_mask.shape[-1] != kv_len:
            valid_mask = None

    if valid_mask is None:
        valid_mask = torch.ones(B, kv_len, device=device, dtype=torch.bool)

    # ---- Decide execution path ----
    num_pages = (kv_len + self.page_size - 1) // self.page_size
    use_paged = self.enable_moe and num_pages > 1 and not is_cross

    if use_paged:
        context, attn_probs = self._paged_moe_attention(q, k_hat, v_hat, valid_mask)
    else:
        # ---- Dense path ----
        scale = 1.0 / math.sqrt(d)
        scores = torch.matmul(q.float(), k_hat.float().transpose(-2, -1).contiguous()) * scale

        if self.is_causal and not is_cross and N > 1:
            diag = 1 + max(0, kv_len - N)
            causal = (
                torch.triu(
                    torch.full(
                        (N, kv_len),
                        float("-inf"),
                        device=device,
                        dtype=torch.float32,
                    ),
                    diagonal=diag,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)
        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)

        attn_probs = F.softmax(scores, dim=-1).to(dtype)
        attn_probs = self.dropout(attn_probs)
        context = torch.matmul(attn_probs, v_hat)

        with torch.no_grad():
            self.aux_loss.zero_()

    # ---- head_mask ----
    if head_mask is not None:
        hm = head_mask.to(context.dtype)
        if hm.dim() == 1:
            hm = hm.view(1, -1, 1, 1)
        elif hm.dim() == 2:
            hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
        elif hm.dim() == 3:
            hm = hm.unsqueeze(1)
        context = context * hm

    # ---- Reshape to (B, N, all_head_size) ----
    context = context[:, :, :N, :] if context.shape[2] != N else context
    context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
    context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

    # ---- Output projection ----
    if self.is_nested_encoder:
        out = context
    else:
        out = self.out_proj(context)

    # MDP-FIX-13: always return a 2-tuple with attn = tensor or None
    attn_ret = attn_probs if output_attentions else None
    return out, attn_ret

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    return {"orig": self._orig_storage_bytes, "comp": self._comp_storage_bytes}

class ModelAdapter (model_or_config: Any,
family: str | None = None)

Expand source code

class ModelAdapter:
    """Read-only architecture contract used by module wrappers.

    Fix MF-05: `family` is normalized to a canonical internal name (e.g. "gpt2"
    becomes "gpt"), while the original name is preserved in `raw_family` for
    logging. Downstream code that inspects `adapter.family` will always see a
    canonical value; code that needs the raw value should use `raw_family`.
    """

    def __init__(self, model_or_config: Any, family: str | None = None):
        self.source = _config(model_or_config)
        raw_family = family or detect_family(model_or_config)
        # Normalize aliases to canonical internal family.
        self.raw_family = raw_family
        self.family = FAMILY_ALIASES.get(raw_family, raw_family)

        self.hidden_size = int(
            getattr(self.source, "hidden_size", getattr(self.source, "d_model", 0)) or 0
        )
        heads = getattr(
            self.source,
            "num_attention_heads",
            getattr(
                self.source,
                "num_heads",
                getattr(self.source, "num_decoder_attention_heads", 0),
            ),
        )
        self.num_heads = int(heads or 0)
        h_dim = getattr(self.source, "head_dim", None)
        self.head_dim = int(
            h_dim
            if h_dim is not None
            else (self.hidden_size // self.num_heads if self.num_heads else 0)
        )
        kv_heads = getattr(
            self.source,
            "num_key_value_heads",
            getattr(
                self.source,
                "num_key_value_heads_decoder",
                getattr(self.source, "num_decoder_attention_heads", self.num_heads or 0),
            ),
        )
        self.num_key_value_heads = int(kv_heads or self.num_heads or 0)
        if self.num_key_value_heads <= 0:
            self.num_key_value_heads = self.num_heads
        self.num_hidden_layers = int(
            getattr(self.source, "num_hidden_layers", getattr(self.source, "num_layers", 0))
        )
        self.q_proj_out = self.num_heads * self.head_dim
        self.kv_proj_out = self.num_key_value_heads * self.head_dim
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_key_value_heads))
        self.projection_aliases = {
            "q": ("q_lin", "q_proj"),
            "k": ("k_lin", "k_proj"),
            "v": ("v_lin", "v_proj"),
            "o": ("out_lin", "o_proj"),
        }

        # ------------------------------------------------------------------
        # Fix MF-01/MF-02/MF-03/MF-04: Consistent causal-family handling.
        # ------------------------------------------------------------------
        # A family is causal if either:
        #   (a) its canonical family is in CAUSAL_FAMILIES, OR
        #   (b) its raw family is in CAUSAL_FAMILIES (defense in depth).
        self.causal = (self.family in CAUSAL_FAMILIES) or (raw_family in CAUSAL_FAMILIES)

        # mask_rank: 4 for causal decoders, 2 for bidirectional encoders, 3 otherwise.
        if self.causal:
            self.mask_rank = 4
        elif self.family in BIDIRECTIONAL_FAMILIES:
            self.mask_rank = 2
        else:
            self.mask_rank = 3

        self.bidirectional = not self.causal

        # position_embedding: rotary for modern decoders, absolute for BERT/GPT-2,
        # relative for T5, none for unknown.
        if self.family in ROTARY_FAMILIES or raw_family in ROTARY_FAMILIES:
            self.position_embedding = "rotary"
        elif (
            self.family in GPT_ABSOLUTE_POSITION_FAMILIES
            or raw_family in GPT_ABSOLUTE_POSITION_FAMILIES
            or self.family in {"bert", "roberta"}
        ):
            self.position_embedding = "absolute"
        elif self.family == "t5":
            self.position_embedding = "relative"
        else:
            self.position_embedding = "none"

        self.supports_gqa = (self.num_key_value_heads < self.num_heads) or self.family == "qwen3_5"
        self.supports_delta_net = self.family == "qwen3_5"
        self.supports_moe = self.family == "qwen3_5"

    def projection(self, module: Any, kind: str) -> Any:
        legacy, native = self.projection_aliases[kind]
        return getattr(module, native, getattr(module, legacy, None))

    def as_dict(self) -> dict[str, Any]:
        return {
            "family": self.family,
            "raw_family": self.raw_family,
            "hidden_size": self.hidden_size,
            "num_heads": self.num_heads,
            "num_key_value_heads": self.num_key_value_heads,
            "q_proj_out": self.q_proj_out,
            "kv_proj_out": self.kv_proj_out,
            "kv_repeat": self.kv_repeat,
            "head_dim": self.head_dim,
            "num_hidden_layers": self.num_hidden_layers,
            "causal": self.causal,
            "mask_rank": self.mask_rank,
            "position_embedding": self.position_embedding,
            "supports_gqa": self.supports_gqa,
            "supports_delta_net": self.supports_delta_net,
            "supports_moe": self.supports_moe,
        }

Read-only architecture contract used by module wrappers.

Fix MF-05: family is normalized to a canonical internal name (e.g. "gpt2" becomes "gpt"), while the original name is preserved in raw_family for logging. Downstream code that inspects adapter.family will always see a canonical value; code that needs the raw value should use raw_family.

Methods

def as_dict(self) ‑> dict[str, typing.Any]

Expand source code

def as_dict(self) -> dict[str, Any]:
    return {
        "family": self.family,
        "raw_family": self.raw_family,
        "hidden_size": self.hidden_size,
        "num_heads": self.num_heads,
        "num_key_value_heads": self.num_key_value_heads,
        "q_proj_out": self.q_proj_out,
        "kv_proj_out": self.kv_proj_out,
        "kv_repeat": self.kv_repeat,
        "head_dim": self.head_dim,
        "num_hidden_layers": self.num_hidden_layers,
        "causal": self.causal,
        "mask_rank": self.mask_rank,
        "position_embedding": self.position_embedding,
        "supports_gqa": self.supports_gqa,
        "supports_delta_net": self.supports_delta_net,
        "supports_moe": self.supports_moe,
    }

def projection(self,
module: Any,
kind: str) ‑> Any

Expand source code

def projection(self, module: Any, kind: str) -> Any:
    legacy, native = self.projection_aliases[kind]
    return getattr(module, native, getattr(module, legacy, None))

class NASCAttention (config: Any,
rank: int | None = 16,
window_size: int = 64,
block_size: int = 32,
top_k: int = 2,
rff_features: int = 64,
pretrained_attn: torch.nn.modules.module.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("nasc", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("nasc_attention", supports_causal=True, supports_gqa=True)
class NASCAttention(BaseAttention):
    r"""NASC: NEXUS-AEGIS Spectral Cascade Attention.
    Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any,
        rank: int | None = 16,
        window_size: int = 64,
        block_size: int = 32,
        top_k: int = 2,
        rff_features: int = 64,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        self.hidden_size: int = getattr(config, "hidden_size", kwargs.get("hidden_size", 768))
        self.num_heads: int = getattr(
            config, "num_attention_heads", kwargs.get("num_attention_heads", 4)
        )

        cfg_head_dim = getattr(config, "head_dim", None) if config else None
        if cfg_head_dim is not None:
            self.head_dim = int(cfg_head_dim)
        elif kwargs.get("head_dim") is not None:
            self.head_dim = int(kwargs["head_dim"])
        elif self.num_heads > 0:
            self.head_dim = self.hidden_size // self.num_heads
        else:
            self.head_dim = 64
        self.all_head_size: int = self.num_heads * self.head_dim

        cfg_kv = (
            getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
            if config
            else None
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        self.window_size: int = window_size
        self.block_size: int = block_size
        self.top_k: int = top_k
        self.rff_features: int = rff_features
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha
        self.enable_cascade_on_encoders: bool = bool(
            kwargs.get(
                "enable_cascade_on_encoders",
                getattr(config, "enable_cascade_on_encoders", True),
            )
        )

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            m_type = str(getattr(config, "model_type", "")).lower() if config else ""
            self.is_decoder = (
                getattr(config, "is_decoder", False)
                or getattr(config, "causal", False)
                or any(
                    m_type.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                if config
                else False
            )

        self.is_nested_encoder: bool = False

        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,
        )

        rff_init = torch.empty(self.head_dim, rff_features)
        nn.init.orthogonal_(rff_init)
        self.rff_proj = nn.Parameter(rff_init, requires_grad=False)
        self.gate_proj = nn.Linear(self.hidden_size, 3)

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    @property
    def query(self) -> LoRALinear:
        return self.q_proj

    @property
    def key(self) -> LoRALinear:
        return self.k_proj

    @property
    def value(self) -> LoRALinear:
        return self.v_proj

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug("UniversalAdapter extraction fallback for %s: %s", proj, e)

        # 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_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            weight = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None))
            if weight is not None:
                weight = weight.data.clone()
                bias = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                bias_data = bias.data.clone() if bias is not None else None
                c_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in c_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size)
                ):
                    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:
            c_proj_mod = src.c_proj
            base_o = getattr(c_proj_mod, "base_layer", c_proj_mod)
            w_o = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None))
            if w_o is not None:
                w_o = w_o.data.clone()
                b_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                p_cls = getattr(base_o, "__class__", type(base_o)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_o, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (w_o.ndim == 2 and w_o.shape[0] < w_o.shape[1])
                ):
                    w_o = w_o.t().contiguous()
                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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj"):
            self.is_nested_encoder = False
        elif (
            (
                hasattr(src, "self")
                and hasattr(src.self, "query")
                and hasattr(src, "output")
                and hasattr(src.output, "dense")
            )
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        # Check required weights
        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("Module copy_weights warning. Proceeding with initialized weights.")

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"NASCAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"NASCAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
                )

        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

        # Copy into projection modules
        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape

                    if proj_key in ("k", "v"):
                        if src_out == out_d:
                            target_w.data.copy_(w_oriented)
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b)
                        elif src_out > out_d and src_out % out_d == 0:
                            rep = src_out // out_d
                            target_w.data.copy_(w_oriented.view(out_d, rep, -1).mean(dim=1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.view(out_d, rep).mean(dim=1))
                        elif src_out < out_d and out_d % src_out == 0:
                            rep = out_d // src_out
                            target_w.data.copy_(w_oriented.repeat(rep, 1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.repeat(rep))
                        else:
                            target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                                w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                            )
                            if b is not None and target.bias is not None:
                                target.bias.data[: min(len(target.bias), len(b))].copy_(
                                    b[: min(len(target.bias), len(b))]
                                )
                    else:
                        target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                            w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                        )
                        if b is not None and target.bias is not None:
                            b_val = b.data if isinstance(b, nn.Parameter) else b
                            target.bias.data[: min(len(target.bias), len(b_val))].copy_(
                                b_val[: min(len(target.bias), len(b_val))]
                            )
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                target.freeze_base_weights()

        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)

    def _local_window_attention(self, Q, K, V, valid_mask, is_causal):
        B, H, L, d = Q.shape
        scale = 1.0 / math.sqrt(d)
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        idx = torch.arange(L, device=Q.device)
        dist = (idx.unsqueeze(0) - idx.unsqueeze(1)).abs()
        mask = dist > self.window_size
        if is_causal:
            # Query i attending to key j: future keys (j > i) must be masked
            mask = mask | (idx.unsqueeze(0) > idx.unsqueeze(1))
        scores = scores.masked_fill(mask.unsqueeze(0).unsqueeze(0), -10000.0)

        # Padding mask
        pad_k = (~valid_mask).unsqueeze(1).unsqueeze(2)
        scores = scores.masked_fill(pad_k, -10000.0)
        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)

    def _sparse_block_attention(self, Q, K, V, valid_mask, is_causal):
        B, H, L, d = Q.shape
        pad_len = (self.block_size - L % self.block_size) % self.block_size
        if pad_len > 0:
            Q_pad = F.pad(Q, (0, 0, 0, pad_len))
            K_pad = F.pad(K, (0, 0, 0, pad_len))
            V_pad = F.pad(V, (0, 0, 0, pad_len))
            valid_pad = F.pad(valid_mask, (0, pad_len), value=False)
        else:
            Q_pad, K_pad, V_pad = Q, K, V
            valid_pad = valid_mask
        Lp = L + pad_len
        nb = Lp // self.block_size

        Qb = Q_pad.view(B, H, nb, self.block_size, d)
        Kb = K_pad.view(B, H, nb, self.block_size, d)
        Vb = V_pad.view(B, H, nb, self.block_size, d)
        valid_b = valid_pad.view(B, nb, self.block_size)

        Q_mean = Qb.mean(dim=3)
        K_mean = Kb.mean(dim=3)
        scale = 1.0 / math.sqrt(d)
        blk_scores = (
            torch.matmul(Q_mean.float(), K_mean.float().transpose(-2, -1).contiguous()) * scale
        )

        if is_causal:
            causal_blk_mask = torch.triu(
                torch.full((nb, nb), float("-inf"), device=Q.device, dtype=torch.float32),
                diagonal=1,
            )
            blk_scores = blk_scores + causal_blk_mask.unsqueeze(0).unsqueeze(0)

        top_k = min(self.top_k, nb)
        _, top_idx = blk_scores.topk(top_k, dim=-1)
        if is_causal:
            for i in range(nb):
                top_idx[:, :, i, 0] = i
        else:
            top_idx[:, :, :, 0] = 0

        out_blocks = torch.zeros_like(Qb)
        for i in range(nb):
            idx = top_idx[:, :, i, :]
            idx_exp = idx.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, -1, self.block_size, d)
            K_sel = torch.gather(Kb, 2, idx_exp).reshape(B, H, top_k * self.block_size, d)
            V_sel = torch.gather(Vb, 2, idx_exp).reshape(B, H, top_k * self.block_size, d)

            scores_i = (
                torch.matmul(Qb[:, :, i].float(), K_sel.float().transpose(-2, -1).contiguous())
                * scale
            )

            if is_causal:
                q_pos = (i * self.block_size + torch.arange(self.block_size, device=Q.device)).view(
                    1, 1, self.block_size, 1
                )
                k_pos = (
                    idx.unsqueeze(-1) * self.block_size
                    + torch.arange(self.block_size, device=Q.device).view(1, 1, 1, self.block_size)
                ).view(B, H, 1, top_k * self.block_size)
                scores_i = scores_i.masked_fill(k_pos > q_pos, -10000.0)

            valid_b_exp = valid_b.unsqueeze(1).expand(B, H, nb, self.block_size)
            idx_v = idx.unsqueeze(-1).expand(-1, -1, -1, self.block_size)
            valid_sel = torch.gather(valid_b_exp, 2, idx_v).reshape(B, H, top_k * self.block_size)
            pad_sel = ~valid_sel.unsqueeze(2)
            scores_i = scores_i.masked_fill(pad_sel, -10000.0)
            scores_i = torch.nan_to_num(scores_i, nan=-1e4, posinf=1e4, neginf=-1e4)

            attn_i = softpick(scores_i, dim=-1).to(Q.dtype)
            out_blocks[:, :, i] = torch.matmul(attn_i, V_sel)

        out = out_blocks.view(B, H, Lp, d)[:, :, :L, :]
        return out

    def _global_spectral_rff_attention(self, Q, K, V, valid_mask, is_causal):
        B, H, L, d = Q.shape
        pad = (~valid_mask).unsqueeze(1).unsqueeze(-1)

        Qf = torch.matmul(Q.float(), self.rff_proj.float())
        Kf = torch.matmul(K.float(), self.rff_proj.float())
        norm = 1.0 / math.sqrt(self.rff_features)
        phi_Q = torch.cat([torch.cos(Qf), torch.sin(Qf)], dim=-1) * norm
        phi_K = torch.cat([torch.cos(Kf), torch.sin(Kf)], dim=-1) * norm

        phi_K = phi_K.masked_fill(pad, 0.0)
        V_masked = V.float().masked_fill(pad, 0.0)

        if is_causal:
            # Causal linear recurrence via cumulative prefix sums
            KV = phi_K.unsqueeze(-1) * V_masked.unsqueeze(-2)  # [B, H, L, 2F, d]
            KV_cumsum = torch.cumsum(KV, dim=2)
            num = torch.matmul(phi_Q.unsqueeze(-2), KV_cumsum).squeeze(-2)  # [B, H, L, d]
            Z = torch.cumsum(phi_K, dim=2)
            den = (phi_Q * Z).sum(dim=-1, keepdim=True) + 1e-5
            return (num / den).to(Q.dtype)
        else:
            KV = torch.matmul(phi_K.transpose(-2, -1).contiguous(), V_masked)
            Z = phi_K.sum(dim=-2, keepdim=True)
            num = torch.matmul(phi_Q, KV)
            den = torch.matmul(phi_Q, Z.transpose(-2, -1).contiguous()) + 1e-5
            return (num / den).to(Q.dtype)

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        L_kv = kv_input.shape[1]

        is_causal = bool(self.is_decoder and not is_cross and L > 1)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is None:
            valid_mask = torch.ones(B, L_kv, device=hidden_states.device, dtype=torch.bool)
        elif eff_mask.dim() == 2:
            valid_mask = eff_mask.bool() if eff_mask.dtype == torch.bool else eff_mask > 0
        elif eff_mask.dim() == 4:
            valid_mask = (
                eff_mask[:, 0, 0, :] > -10000.0
                if eff_mask.min() < -1000
                else eff_mask[:, 0, 0, :] > 0
            )
        else:
            valid_mask = torch.ones(B, L_kv, device=hidden_states.device, dtype=torch.bool)

        Q = self.q_proj(hidden_states).view(B, L, H, d).permute(0, 2, 1, 3)
        K = self.k_proj(kv_input).view(B, L_kv, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, L_kv, 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 and not is_cross:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        use_dense = bool(is_cross or (not self.is_decoder and not self.enable_cascade_on_encoders))

        if not hasattr(self, "_logged_once"):
            logger.debug(
                "NASCAttention: L=%d, L_kv=%d, is_causal=%s, head_dim=%d, all_head_size=%d, path=%s",
                L,
                L_kv,
                is_causal,
                self.head_dim,
                self.all_head_size,
                "dense" if use_dense else "cascade",
            )
            self._logged_once = True

        if use_dense:
            # Fall back to standard scaled dot product attention for short/cross sequences or encoders
            scale = 1.0 / math.sqrt(d)
            scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
            if is_causal:
                causal_mask = (
                    torch.triu(
                        torch.full(
                            (L, L_kv),
                            float("-inf"),
                            device=hidden_states.device,
                            dtype=torch.float32,
                        ),
                        diagonal=1,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal_mask

            if eff_mask is not None:
                if eff_mask.dim() == 2:
                    scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
                elif eff_mask.dim() == 3:
                    scores = scores + eff_mask.unsqueeze(1).float()
                elif eff_mask.dim() == 4:
                    scores = scores + eff_mask.float()

            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)
            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)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(-1)
                attn = attn * head_mask
            O = torch.matmul(attn, V)
        else:
            summary = hidden_states.mean(dim=1)
            gates = F.softmax(self.gate_proj(summary), dim=-1)  # [B, 3]
            g0 = gates[:, 0].view(B, 1, 1, 1)
            g1 = gates[:, 1].view(B, 1, 1, 1)
            g2 = gates[:, 2].view(B, 1, 1, 1)

            O_loc = self._local_window_attention(Q, K, V, valid_mask, is_causal)
            O_sp = self._sparse_block_attention(Q, K, V, valid_mask, is_causal)
            O_glb = self._global_spectral_rff_attention(Q, K, V, valid_mask, is_causal)

            O = g0 * O_loc + g1 * O_sp + g2 * O_glb

            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.unsqueeze(-1)
                else:
                    hm = head_mask
                O = O * hm

        context = O.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_proj(context)

        return out, None

NASC: NEXUS-AEGIS Spectral Cascade Attention. Hardened for multi-architecture deployment across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.

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.

Instance variables

prop key : LoRALinear

Expand source code

@property
def key(self) -> LoRALinear:
    return self.k_proj

prop query : LoRALinear

Expand source code

@property
def query(self) -> LoRALinear:
    return self.q_proj

prop value : LoRALinear

Expand source code

@property
def value(self) -> LoRALinear:
    return self.v_proj

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class NEXUSAttention (config: Any | None = None,
model_config: Any | None = None,
nexus_cfg: NEXUSConfig | None = None,
layer_idx: int = 0,
total_layers: int = 6,
pretrained_attn: torch.nn.modules.module.Module | None = None,
bert_mode: bool = False,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("nexus", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("nexus_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("NEXUSAttention", supports_causal=True, supports_gqa=True)
class NEXUSAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        nexus_cfg: NEXUSConfig | None = None,
        layer_idx: int = 0,
        total_layers: int = 6,
        pretrained_attn: nn.Module | None = None,
        bert_mode: bool = False,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        if model_config is None:
            model_config = config
        super().__init__(model_config)
        self.num_heads: int = getattr(
            model_config,
            "num_attention_heads",
            getattr(model_config, "num_heads", kwargs.get("num_attention_heads", 4)),
        )
        self.hidden_size: int = getattr(model_config, "hidden_size", kwargs.get("hidden_size", 768))
        cfg_head_dim = getattr(model_config, "head_dim", None) if model_config else None
        adapter_head_dim = (
            getattr(getattr(model_config, "adapter", None), "head_dim", None)
            if model_config
            else None
        )
        self.head_dim: int = int(
            cfg_head_dim
            or adapter_head_dim
            or kwargs.get("head_dim")
            or (self.hidden_size // self.num_heads if self.num_heads > 0 else 64)
        )
        self.all_head_size: int = self.num_heads * self.head_dim

        cfg_kv = (
            getattr(
                model_config,
                "num_key_value_heads",
                getattr(model_config, "num_kv_heads", None),
            )
            if model_config
            else None
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        self.cfg: NEXUSConfig = nexus_cfg if nexus_cfg is not None else NEXUSConfig()
        self.layer_idx: int = layer_idx
        self.bert_mode: bool = bert_mode
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(getattr(model_config, "model_type", "") or "").lower() if model_config else ""
            self.is_decoder = (
                bool(
                    getattr(model_config, "is_decoder", False)
                    or getattr(model_config, "causal", False)
                )
                or any(
                    _mt.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "falcon",
                        "bloom",
                        "opt",
                    )
                )
                if model_config
                else False
            )

        self.is_nested_encoder: bool = False

        src_p = next(pretrained_attn.parameters(), None) if pretrained_attn is not None else None
        target_dtype = kwargs.get("dtype", src_p.dtype if src_p is not None else torch.float32)
        target_device = kwargs.get("device", src_p.device if src_p is not None else None)

        self.query = LoRALinear(
            self.hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            device=target_device,
            dtype=target_dtype,
        )
        self.key = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            device=target_device,
            dtype=target_dtype,
        )
        self.value = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            device=target_device,
            dtype=target_dtype,
        )
        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,
            device=target_device,
            dtype=target_dtype,
        )
        self.dropout = nn.Dropout(self.cfg.dropout)

        # NX-04: pass is_decoder into DSS
        self.dss_attn = DualStageStripeAttention(
            self.cfg,
            self.head_dim,
            num_heads=self.num_heads,
            is_decoder=self.is_decoder,
            is_nested_encoder=self.is_nested_encoder,
        )

        self.seer_gate: SeerGating | None = None
        if self.cfg.seer_enabled:
            self.seer_gate = SeerGating(
                self.head_dim,
                self.num_heads,
                self.cfg.seer_gate_rank,
                self.cfg.seer_top_k_blocks,
                self.cfg.seer_temperature,
            )

        self.rect_interval = self.cfg.rectify_interval
        self.is_rectify = (
            (self.rect_interval > 0)
            and (layer_idx % self.rect_interval == 0)
            and ((pretrained_attn is None) or getattr(self.cfg, "enable_spectraformer", False))
        )
        if self.is_rectify:
            self.spec_rectify: SpectraformerRectify | None = SpectraformerRectify(
                self.head_dim, self.cfg.rff_features
            )
        else:
            self.spec_rectify = None

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

        if target_device is not None and target_dtype is not None:
            self.to(device=target_device, dtype=target_dtype)
        elif target_device is not None:
            self.to(device=target_device)
        elif target_dtype is not None:
            self.to(target_dtype)

    @property
    def q_proj(self):
        return self.query

    @property
    def k_proj(self):
        return self.key

    @property
    def v_proj(self):
        return self.value

    @property
    def q_lin(self):
        return self.query

    @property
    def k_lin(self):
        return self.key

    @property
    def v_lin(self):
        return self.value

    @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

    def _reshape(self, x: torch.Tensor, n_heads: int | None = None) -> torch.Tensor:
        B, N, _ = x.shape
        heads = n_heads if n_heads is not None else self.num_heads
        return x.view(B, N, heads, self.head_dim).transpose(1, 2).contiguous()

    # NX-05: robust valid-mask extraction
    def _get_token_valid_mask(self, mask, B, N, device):
        if mask is None:
            return None
        if mask.dim() == 2:
            if mask.dtype == torch.bool:
                return mask.to(device)
            if mask.dtype.is_floating_point:
                return (mask > -1e4).to(device)
            return (mask > 0).to(device)
        if mask.dim() == 3:
            if mask.dtype == torch.bool:
                v = mask
            elif mask.dtype.is_floating_point:
                v = mask > -1e4
            else:
                v = mask > 0
            return v.any(dim=1).to(device)
        if mask.dim() == 4:
            if mask.dtype == torch.bool:
                v = mask
            elif mask.dtype.is_floating_point:
                v = mask > -1e4
            else:
                v = mask > 0
            return v.any(dim=-2).squeeze(1).to(device)
        return torch.ones(B, N, device=device, dtype=torch.bool)

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        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

        if hasattr(src, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            weight = src.c_attn.weight.data
            bias = getattr(src.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                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:
            w_o = src.c_proj.weight.data
            if (
                w_o.dim() == 2
                and w_o.shape[0] != self.hidden_size
                and w_o.shape[1] == self.hidden_size
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(src.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        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))

        named_map = {
            "q": ["q_proj", "q_lin", "query"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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

        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        is_nested = False
        if (
            hasattr(src, "self")
            and hasattr(src.self, "query")
            or (
                cls_name
                in (
                    "BertSelfAttention",
                    "RobertaSelfAttention",
                    "XLMRobertaSelfAttention",
                )
                or (cls_name.endswith("SelfAttention") and not has_out_proj)
            )
            and not has_out_proj
        ):
            is_nested = True
        self.is_nested_encoder = bool(is_nested)
        # NX-09: propagate to DSS
        if hasattr(self, "dss_attn"):
            self.dss_attn.is_nested_encoder = self.is_nested_encoder

        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("Module copy_weights warning. Proceeding with initialized weights.")

        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(f"NEXUSAttention extracted weight for {proj_key} near-zero norm")

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning("NEXUSAttention output weight near-zero norm")

        def _orient(w: torch.Tensor, in_dim: int) -> torch.Tensor:
            if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim:
                return w.t().contiguous()
            return w

        with torch.no_grad():
            for proj_key, target in [
                ("q", self.query),
                ("k", self.key),
                ("v", self.value),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape

                    if proj_key in ("k", "v"):
                        if src_out == out_d:
                            target_w.data.copy_(w_oriented)
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b)
                        elif src_out > out_d and src_out % out_d == 0:
                            rep = src_out // out_d
                            target_w.data.copy_(w_oriented.view(out_d, rep, -1).mean(dim=1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.view(out_d, rep).mean(dim=1))
                        elif src_out < out_d and out_d % src_out == 0:
                            rep = out_d // src_out
                            target_w.data.copy_(w_oriented.repeat(rep, 1))
                            if b is not None and target.bias is not None:
                                target.bias.data.copy_(b.repeat(rep))
                        else:
                            target_w.data.zero_()
                            target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                                w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                            )
                            if b is not None and target.bias is not None:
                                target.bias.data.zero_()
                                target.bias.data[: min(len(target.bias), len(b))].copy_(
                                    b[: min(len(target.bias), len(b))]
                                )
                    else:
                        target_w.data.zero_()
                        target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                            w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                        )
                        if b is not None and target.bias is not None:
                            b_val = b.data if isinstance(b, nn.Parameter) else b
                            target.bias.data.zero_()
                            target.bias.data[: min(len(target.bias), len(b_val))].copy_(
                                b_val[: min(len(target.bias), len(b_val))]
                            )
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                target.freeze_base_weights()

        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)
            if hasattr(self, "dss_attn"):
                self.dss_attn.is_decoder = self.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,
        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
        H, D = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        N_kv = kv_input.shape[1]

        Q = self._reshape(self.query(hidden_states), self.num_heads)
        K = self._reshape(self.key(kv_input), self.num_kv_heads)
        V = self._reshape(self.value(kv_input), self.num_kv_heads)

        # NX-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 = 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]
            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.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is None:
            eff_mask = kwargs.get("attn_mask")
        token_valid = self._get_token_valid_mask(eff_mask, B, N_kv, device)

        is_causal = bool(self.is_decoder and not is_cross)
        min_seq_len = getattr(self.cfg, "sparse_min_seq_len", 8)
        use_dense = is_causal or is_cross or (min_seq_len >= N) or (self.dss_attn is None)

        if use_dense:
            scale = 1.0 / math.sqrt(D)
            scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
            if is_causal:
                causal_diag = 1 + max(0, N_kv - N)
                causal_mask = (
                    torch.triu(
                        torch.full((N, N_kv), float("-inf"), device=device, dtype=torch.float32),
                        diagonal=causal_diag,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal_mask

            if eff_mask is not None:
                if eff_mask.dim() == 2:
                    scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
                elif eff_mask.dim() == 3:
                    scores = scores + eff_mask.unsqueeze(1).float()
                elif eff_mask.dim() == 4:
                    scores = scores + eff_mask.float()

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            attn_weights = F.softmax(scores, dim=-1).to(dtype)
            attn_weights = self.dropout(attn_weights)

            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)
                elif head_mask.dim() == 3:
                    head_mask = head_mask[:, :, :1].unsqueeze(-1)
                attn_weights = attn_weights * head_mask

            context = torch.matmul(attn_weights, V)
        else:
            if self.is_rectify and self.spec_rectify is not None:
                context = self.spec_rectify(Q, K, V)
            else:
                seer_mask = None
                if self.seer_gate is not None:
                    bs = self.cfg.block_size
                    nb = (N + bs - 1) // bs
                    seer_mask = self.seer_gate(Q, K, nb, training=self.training)

                context = self.dss_attn(Q, K, V, token_valid, seer_mask=seer_mask)

            if token_valid is not None:
                context = context * token_valid.unsqueeze(1).unsqueeze(-1).to(context.dtype)

        context = context.transpose(1, 2).contiguous().view(B, N, H * D)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(self.dropout(context))

        return (output, None)

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.tensor(0.0, device=self.query.weight.device)

Every 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.

Instance variables

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.key

prop k_proj

Expand source code

@property
def k_proj(self):
    return self.key

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.query

prop q_proj

Expand source code

@property
def q_proj(self):
    return self.query

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.value

prop v_proj

Expand source code

@property
def v_proj(self):
    return self.value

Methods

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.tensor(0.0, device=self.query.weight.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class NEXUSConfig (block_size: int = 16,
top_k_blocks: int = 4,
xattn_enabled: bool = True,
xattn_stride: int = 2,
seer_enabled: bool = True,
seer_gate_rank: int = 16,
seer_top_k_blocks: int = 4,
seer_temperature: float = 1.0,
dual_stage_enabled: bool = True,
intra_block_full: bool = True,
stripe_enabled: bool = True,
stripe_width: int = 4,
rectify_interval: int = 4,
rff_features: int = 128,
softpick_enabled: bool = True,
softpick_temperature: float = 1.0,
sparse_min_seq_len: int = 8,
enable_spectraformer: bool = False,
dropout: float = 0.1)

Expand source code

@dataclass
class NEXUSConfig:
    block_size: int = 16
    top_k_blocks: int = 4
    xattn_enabled: bool = True
    xattn_stride: int = 2
    seer_enabled: bool = True
    seer_gate_rank: int = 16
    seer_top_k_blocks: int = 4
    seer_temperature: float = 1.0
    dual_stage_enabled: bool = True
    intra_block_full: bool = True
    stripe_enabled: bool = True
    stripe_width: int = 4
    rectify_interval: int = 4
    rff_features: int = 128
    softpick_enabled: bool = True
    softpick_temperature: float = 1.0
    sparse_min_seq_len: int = 8
    enable_spectraformer: bool = False
    dropout: float = 0.1

NEXUSConfig(block_size: int = 16, top_k_blocks: int = 4, xattn_enabled: bool = True, xattn_stride: int = 2, seer_enabled: bool = True, seer_gate_rank: int = 16, seer_top_k_blocks: int = 4, seer_temperature: float = 1.0, dual_stage_enabled: bool = True, intra_block_full: bool = True, stripe_enabled: bool = True, stripe_width: int = 4, rectify_interval: int = 4, rff_features: int = 128, softpick_enabled: bool = True, softpick_temperature: float = 1.0, sparse_min_seq_len: int = 8, enable_spectraformer: bool = False, dropout: float = 0.1)

Instance variables

var block_size : int

The type of the None singleton.

var dropout : float

The type of the None singleton.

var dual_stage_enabled : bool

The type of the None singleton.

var enable_spectraformer : bool

The type of the None singleton.

var intra_block_full : bool

The type of the None singleton.

var rectify_interval : int

The type of the None singleton.

var rff_features : int

The type of the None singleton.

var seer_enabled : bool

The type of the None singleton.

var seer_gate_rank : int

The type of the None singleton.

var seer_temperature : float

The type of the None singleton.

var seer_top_k_blocks : int

The type of the None singleton.

var softpick_enabled : bool

The type of the None singleton.

var softpick_temperature : float

The type of the None singleton.

var sparse_min_seq_len : int

The type of the None singleton.

var stripe_enabled : bool

The type of the None singleton.

var stripe_width : int

The type of the None singleton.

var top_k_blocks : int

The type of the None singleton.

var xattn_enabled : bool

The type of the None singleton.

var xattn_stride : int

The type of the None singleton.

class NaiveSVDAttention (config: Any | None = None,
model_config: Any | None = None,
energy: float = 1.0,
dropout: float = 0.1,
pretrained_attn: nn.Module | None = None,
rank: int | None = None,
lora_rank: int = 0,
lora_alpha: float = 1.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("naive_svd", supports_causal=True, supports_gqa=True)
class NaiveSVDAttention(MinimalSVDAttention):
    r"""
    Naive truncated SVD baseline attention module.
    Inherits production-hardened multi-architecture routing from MinimalSVDAttention
    with uncalibrated plain truncated SVD initialization and no recovery adapters.
    """

    is_custom_attention: bool = True
    supports_causal: bool = True
    supports_gqa: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        energy: float = 1.0,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        rank: int | None = None,
        lora_rank: int = 0,  # Zero recovery parameters for naive baseline
        lora_alpha: float = 1.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            config=config,
            model_config=model_config,
            energy=energy,
            dropout=dropout,
            pretrained_attn=pretrained_attn,
            rank=rank,
            lora_rank=0,
            lora_alpha=1.0,
            is_decoder=is_decoder,
            **kwargs,
        )

Naive truncated SVD baseline attention module. Inherits production-hardened multi-architecture routing from MinimalSVDAttention with uncalibrated plain truncated SVD initialization and no recovery adapters.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• MinimalSVDAttention
• BaseAttention
• torch.nn.modules.module.Module

Inherited members

• MinimalSVDAttention:
  • copy_weights
  • extra_repr
  • forward
  • is_custom_attention
  • supports_causal
  • supports_gqa

class Namespace (**kwargs)

Expand source code

class Namespace(_AttributeHolder):
    """Simple object for storing attributes.

    Implements equality by attribute names and values, and provides a simple
    string representation.
    """

    def __init__(self, **kwargs):
        for name in kwargs:
            setattr(self, name, kwargs[name])

    def __eq__(self, other):
        if not isinstance(other, Namespace):
            return NotImplemented
        return vars(self) == vars(other)

    def __contains__(self, key):
        return key in self.__dict__

Simple object for storing attributes.

Implements equality by attribute names and values, and provides a simple string representation.

Ancestors

• argparse._AttributeHolder

class NativeHybridAttention (config: Any | None = None,
rank: int = 64,
window_size: int = 64,
rff_features: int = 128,
dropout: float = 0.1,
pretrained_attn: nn.Module | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("nha_class")
@ATTENTION_REGISTRY.register("nha_attention_class")
class NativeHybridAttention(BaseAttention):
    """
    Native Hybrid Attention: local sliding-window + global RFF, fused by a
    token-level learned gate.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 64,
        window_size: int = 64,
        rff_features: int = 128,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        self.config = config
        # NHA-FIX-12: robust comp_config resolution
        if comp_config is None:
            comp_config = getattr(config, "compression_config", None)
        self.comp_config = comp_config

        source_attn = (
            pretrained_attn
            if pretrained_attn is not None
            else kwargs.get("source_attn", kwargs.get("old_attn"))
        )

        # ---------------- dimension resolution ----------------
        hidden_size = kwargs.get(
            "hidden_size",
            kwargs.get("dim", kwargs.get("d_model")),
        )
        num_heads = kwargs.get(
            "num_attention_heads",
            kwargs.get("num_heads", kwargs.get("n_heads")),
        )
        num_kv_heads = kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads"))

        if hidden_size is None or num_heads is None:
            if config is not None:
                if hidden_size is None:
                    hidden_size = _cfg_get(
                        config,
                        "hidden_size",
                        _cfg_get(config, "dim", _cfg_get(config, "d_model", 768)),
                    )
                if num_heads is None:
                    num_heads = _cfg_get(
                        config,
                        "num_attention_heads",
                        _cfg_get(config, "n_heads", _cfg_get(config, "num_heads", 12)),
                    )
                if num_kv_heads is None:
                    num_kv_heads = _cfg_get(
                        config,
                        "num_key_value_heads",
                        _cfg_get(config, "num_kv_heads", num_heads),
                    )
            elif source_attn is not None:
                if hidden_size is None:
                    out_layer = getattr(
                        source_attn,
                        "out_lin",
                        getattr(
                            source_attn,
                            "dense",
                            getattr(
                                source_attn,
                                "c_proj",
                                getattr(source_attn, "o_proj", None),
                            ),
                        ),
                    )
                    hidden_size = getattr(
                        out_layer,
                        "out_features",
                        getattr(
                            source_attn,
                            "hidden_size",
                            getattr(source_attn, "embed_dim", 768),
                        ),
                    )
                if num_heads is None:
                    num_heads = getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(
                            source_attn,
                            "n_heads",
                            getattr(source_attn, "num_heads", 12),
                        ),
                    )
                if num_kv_heads is None:
                    num_kv_heads = getattr(
                        source_attn,
                        "num_key_value_heads",
                        getattr(source_attn, "num_kv_heads", num_heads),
                    )
            else:
                hidden_size = 768
                num_heads = 12
                num_kv_heads = num_heads

        if num_kv_heads is None:
            num_kv_heads = num_heads

        self.dim = int(hidden_size)
        self.hidden_size = self.dim
        self.num_heads = int(num_heads)
        # NHA-FIX-06: clamp num_kv_heads
        self.num_kv_heads = min(int(num_kv_heads), self.num_heads)
        self.head_dim = self.dim // self.num_heads if self.num_heads > 0 else 64
        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.window_size = int(
            _cfg_get(
                comp_config,
                "window_size",
                _cfg_get(config, "window_size", kwargs.get("window_size", window_size)),
            )
        )
        self.rff_features = int(
            _cfg_get(
                comp_config,
                "rff_features",
                _cfg_get(config, "rff_features", kwargs.get("rff_features", rff_features)),
            )
        )

        # ---------------- decoder detection ----------------
        model_type = str(getattr(config, "model_type", "") or "").lower() if config else ""
        is_dec = (
            getattr(
                config,
                "is_decoder",
                getattr(config, "is_causal", kwargs.get("is_causal", False)),
            )
            if config
            else False
        )
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if any(model_type.startswith(p) for p in decoder_prefixes):
            is_dec = True
        elif hasattr(config, "architectures"):
            archs = getattr(config, "architectures", []) or []
            if any("GPT" in str(a) or "CausalLM" in str(a) for a in archs):
                is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        # NHA-FIX-07
        self.is_nested_encoder = False
        self.output_scale = nn.Parameter(torch.ones(1))

        # ---------------- LoRA config ----------------
        def _get_lora_rank():
            if config is not None:
                v = getattr(config, "lora_rank", None)
                if v is not None:
                    return int(v)
            if comp_config is not None:
                v = getattr(comp_config, "lora_rank", None)
                if v is not None:
                    return int(v)
            return int(kwargs.get("lora_rank", 16))

        def _get_lora_alpha():
            if config is not None:
                v = getattr(config, "lora_alpha", None)
                if v is not None:
                    return float(v)
            if comp_config is not None:
                v = getattr(comp_config, "lora_alpha", None)
                if v is not None:
                    return float(v)
            return float(kwargs.get("lora_alpha", 32.0))

        self.lora_rank = _get_lora_rank()
        self.lora_alpha = _get_lora_alpha()
        self.freeze_base = self.lora_rank > 0

        # ---------------- projections ----------------
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Aliases
        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj
        self.q_lin = self.q_proj
        self.k_lin = self.k_proj
        self.v_lin = self.v_proj
        self.out_lin = self.out_proj
        self.dense = self.out_proj
        self.c_proj = self.out_proj
        self.o_proj = self.out_proj

        omega = self._build_orthogonal_matrix(self.head_dim, self.rff_features)
        self.register_buffer("rff_proj", omega)

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(config, "attention_probs_dropout_prob", dropout),
            )
            if config
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)

        self.gate = nn.Linear(self.dim, 1)
        # Bias strongly towards local at init: sigmoid(3.0) ≈ 0.95
        nn.init.zeros_(self.gate.weight)
        nn.init.constant_(self.gate.bias, 3.0)

        self.rotary_emb = (
            getattr(source_attn, "rotary_emb", None) if source_attn is not None else None
        )

        if source_attn is not None:
            self.copy_weights(source_attn)

        # NHA-FIX-13: forward sanity check
        try:
            with torch.no_grad():
                p = next(self.parameters())
                dummy = torch.randn(1, 4, self.dim, device=p.device, dtype=p.dtype)
                _ = self.forward(dummy)
        except Exception as e:
            logger.debug("NHA forward sanity check warning: %s", e)

    # ------------------------------------------------------------------
    def _build_orthogonal_matrix(self, d: int, m: int) -> torch.Tensor:
        blocks = []
        remaining = m
        while remaining > 0:
            block_size = min(d, remaining)
            g = torch.randn(d, block_size)
            q, _ = torch.linalg.qr(g)
            blocks.append(q)
            remaining -= block_size
        W = torch.cat(blocks, dim=1) * math.sqrt(d)
        return W

    # ------------------------------------------------------------------
    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
        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[0] < c_w.shape[1])
                )
                if is_conv:
                    c_w = c_w.t().contiguous()
                split = 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, :],
                    c_b[:split] if c_b is not None else None,
                )
                extracted["k"] = (
                    c_w[split : 2 * split, :],
                    c_b[split : 2 * split] if c_b is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :, :],
                    c_b[2 * split :] 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. RoBERTa nested
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for 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[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
        else:
            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 key, names in candidates.items():
                for n in names:
                    if hasattr(src, n):
                        mod = getattr(src, n)
                        if mod is not None and hasattr(mod, "weight"):
                            extracted[key] = (
                                mod.weight.data.clone(),
                                mod.bias.data.clone() if mod.bias is not None else None,
                            )
                            break

        # 4. UniversalAdapter fallback
        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:
                # NHA-FIX-03: no print
                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"NativeHybridAttention.copy_weights: missing mandatory "
                f"projections {missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for key, target_mod in target_map.items():
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    else:
                        raise RuntimeError(f"NHA: missing weights for {key}")

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = float(torch.linalg.norm(target_mod.weight.float()).item())
                if fnorm < 1e-4:
                    logger.warning("NHA: '%s' weight norm too small (%.2e).", key, fnorm)

        if self.lora_rank > 0:
            self.q_proj.freeze_base_weights()
            self.k_proj.freeze_base_weights()
            self.v_proj.freeze_base_weights()
            self.out_proj.freeze_base_weights()

    # ------------------------------------------------------------------
    def _local_window_attention(self, Q, K, V, valid_mask, is_causal=False):
        B, H, N, d = Q.shape
        scale = 1.0 / math.sqrt(d)
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1)) * scale

        idx = torch.arange(N, device=Q.device)
        i = idx.unsqueeze(1)
        j = idx.unsqueeze(0)

        if is_causal:
            band_mask = (j > i) | ((i - j) > self.window_size)
        else:
            band_mask = (i - j).abs() > self.window_size

        scores = scores.masked_fill(band_mask.unsqueeze(0).unsqueeze(0), -1e4)
        pad_mask = ~valid_mask.unsqueeze(1).unsqueeze(2)
        scores = scores.masked_fill(pad_mask, -1e4)

        attn = F.softmax(scores, dim=-1).to(Q.dtype)
        attn = self.dropout(attn)
        return torch.matmul(attn, V)

    # ------------------------------------------------------------------
    def _global_rff_attention(self, Q, K, V, valid_mask, is_causal=False):
        B, H, N, d = Q.shape
        scale = d**-0.25
        Q_scaled = Q * scale
        K_scaled = K * scale

        proj = self.rff_proj.to(dtype=Q.dtype, device=Q.device)
        Qf = torch.matmul(Q_scaled, proj)
        Kf = torch.matmul(K_scaled, proj)

        norm = 1.0 / math.sqrt(self.rff_features)
        phi_Q = (F.elu(Qf) + 1.0) * norm
        phi_K = (F.elu(Kf) + 1.0) * norm

        pad = valid_mask.unsqueeze(1).unsqueeze(-1).to(Q.dtype)
        phi_K = phi_K * pad
        V_masked = V * pad

        if is_causal:
            kv_outer = phi_K.unsqueeze(-1) * V_masked.unsqueeze(-2)
            cum_KV = torch.cumsum(kv_outer.float(), dim=2).to(Q.dtype)
            num = (phi_Q.unsqueeze(-2) @ cum_KV).squeeze(-2)

            cum_K = torch.cumsum(phi_K.float(), dim=2).to(Q.dtype)
            den = (phi_Q * cum_K).sum(dim=-1, keepdim=True).clamp(min=1e-5)
            out = (num / den) * pad
        else:
            phi_Q = phi_Q * pad
            KV = torch.matmul(phi_K.transpose(-2, -1), V_masked)
            Z = phi_K.sum(dim=-2, keepdim=True)
            num = torch.matmul(phi_Q, KV)
            den = torch.matmul(phi_Q, Z.transpose(-2, -1)).clamp(min=1e-5)
            out = (num / den) * pad

        return out

    # ------------------------------------------------------------------
    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]:
        # NHA-FIX-08: explicit error instead of silent bypass
        if encoder_hidden_states is not None:
            raise NotImplementedError("NHA does not support cross-attention.")

        while isinstance(hidden_states, (tuple, list)):
            hidden_states = hidden_states[0]

        B, N, D = hidden_states.shape
        is_causal = self.is_decoder or kwargs.get("is_causal", False)
        valid_mask = get_token_valid_mask(attention_mask, B, N, hidden_states.device)

        Q = self.q_proj(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        K_new = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        V_new = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )

        # ---------------- RoPE ----------------
        if position_embeddings is None:
            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_new = apply_rotary_pos_emb(Q, K_new, cos, sin, kwargs.get("position_ids"))
        elif self.rotary_emb is not None:
            pos_ids = kwargs.get("position_ids")
            if pos_ids is None:
                pos_ids = torch.arange(N, device=hidden_states.device).unsqueeze(0).expand(B, -1)
            try:
                # NHA-FIX-11: pass a meaningful reference tensor
                cos, sin = self.rotary_emb(K_new, pos_ids)
                Q, K_new = apply_rotary_pos_emb(Q, K_new, cos, sin, pos_ids)
            except Exception as e:
                # NHA-FIX-02: do not silently return None
                raise RuntimeError(f"NHA RoPE failed: {e}") from e

        # ---------------- KV cache (NHA-FIX-05) ----------------
        past_k = past_v = None
        past_len = 0
        if past_key_value is not None:
            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]
                past_len = past_k.shape[-2]
            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]
                past_len = past_k.shape[-2]

        if past_k is not None:
            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"):
            past_key_value.update(K_new, V_new, layer_idx)

        K_cat.shape[-2]

        # GQA
        if self.num_heads != self.num_kv_heads:
            num_reps = self.num_heads // self.num_kv_heads
            K_full = repeat_kv(K_cat, num_reps)
            V_full = repeat_kv(V_cat, num_reps)
        else:
            K_full = K_cat
            V_full = V_cat

        # Combined valid mask (past + current): assume past is valid
        if past_len > 0:
            all_valid = torch.cat(
                [
                    torch.ones(B, past_len, device=hidden_states.device, dtype=torch.bool),
                    valid_mask,
                ],
                dim=1,
            )
        else:
            all_valid = valid_mask

        local_out = self._local_window_attention(Q, K_full, V_full, all_valid, is_causal=is_causal)
        global_out = self._global_rff_attention(Q, K_full, V_full, all_valid, is_causal=is_causal)

        gate = torch.sigmoid(self.gate(hidden_states)).unsqueeze(1)  # (B, 1, N, 1)

        out = gate * local_out + (1.0 - gate) * global_out
        out = out * valid_mask.unsqueeze(1).unsqueeze(-1).to(out.dtype)

        if head_mask is not None:
            if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                out = out * head_mask.view(1, self.num_heads, 1, 1).to(out.dtype)
            elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                out = out * head_mask.view(head_mask.size(0), self.num_heads, 1, 1).to(out.dtype)
            elif head_mask.dim() == 4:
                out = out * head_mask.to(out.dtype)

        out = out.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)

        # NHA-FIX-15: output_scale applied consistently
        if self.is_nested_encoder:
            out = out * self.output_scale
            return out, None

        out = out * self.output_scale
        out = self.out_proj(out)

        # NHA-FIX-09: no silent pad/truncate
        if out.shape[-1] != D:
            raise RuntimeError(
                f"NHA: output dim {out.shape[-1]} != input dim {D}. "
                "This indicates a projection configuration bug."
            )

        return out, None

Native Hybrid Attention: local sliding-window + global RFF, fused by a token-level learned gate.

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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class NSAAttention (config: Any | None = None,
rank: int | None = None,
window_size: int = 32,
top_k: int = 16,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "nsa", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "nsa_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class NativeSparseAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        window_size: int = 32,
        top_k: int = 16,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        hidden_size = kwargs.get("hidden_size")
        num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")
        num_kv_heads = kwargs.get("num_key_value_heads") or kwargs.get("num_kv_heads")

        if hidden_size is None:
            if config is not None:
                hidden_size = getattr(config, "hidden_size", getattr(config, "dim", 768))
            elif source_attn is not None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                hidden_size = (
                    getattr(out_layer, "out_features", 768) if out_layer is not None else 768
                )
            else:
                hidden_size = 768

        if num_heads is None:
            if config is not None:
                num_heads = getattr(config, "num_attention_heads", getattr(config, "n_heads", 12))
            elif source_attn is not None:
                num_heads = getattr(
                    source_attn,
                    "num_attention_heads",
                    getattr(source_attn, "n_heads", 12),
                )
            else:
                num_heads = 12

        if num_kv_heads is None:
            if config is not None:
                num_kv_heads = getattr(
                    config,
                    "num_key_value_heads",
                    getattr(config, "num_kv_heads", num_heads),
                )
            elif source_attn is not None:
                num_kv_heads = getattr(source_attn, "num_key_value_heads", num_heads)
            else:
                num_kv_heads = 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 = self.dim // self.num_heads if self.num_heads > 0 else 64
        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.window_size = (
            getattr(config, "nsa_window_size", window_size) if config else window_size
        )
        self.top_k = getattr(config, "nsa_top_k", top_k) if config else top_k

        # Decoder detection
        model_type = getattr(config, "model_type", "").lower() if config else ""
        is_dec = (
            getattr(
                config,
                "is_decoder",
                getattr(config, "is_causal", kwargs.get("is_causal", False)),
            )
            if config
            else False
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(_mt.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        # LoRA config
        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0

        # Projections
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        if source_attn is not None:
            self.copy_weights(source_attn)

    def copy_weights(self, src: Any) -> None:
        if src is None:
            return

        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")
        )
        if cls_name in ["BertSelfAttention", "RobertaSelfAttention"] or (
            hasattr(src, "query") and not has_out
        ):
            self.is_nested_encoder = True
        if hasattr(src, "self") and hasattr(src, "output"):
            self.is_nested_encoder = True

        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
        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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (c_proj_mod.__class__.__name__ == "Conv1D")
                        or (base_o.__class__.__name__ == "Conv1D")
                        or hasattr(c_proj_mod, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        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. 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
        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"):
                            w_data = mod.weight.data.clone()
                            if (
                                (mod.__class__.__name__ == "Conv1D")
                                or hasattr(mod, "nf")
                                or (
                                    w_data.ndim == 2
                                    and proj_key == "o"
                                    and w_data.shape[0] < w_data.shape[1]
                                )
                            ):
                                w_data = w_data.t().contiguous()
                            extracted[proj_key] = (
                                w_data,
                                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(f"UniversalAdapter fallback: {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"NativeSparseAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                w_norm = float(torch.norm(target_mod.weight.data).item())
                if w_norm < 1e-4:
                    logger.warning(
                        f"NativeSparseAttention weight verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                    )

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    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,
        output_attentions: bool | None = False,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = self.q_proj(hidden_states).view(B, L, H, d).permute(0, 2, 1, 3)
        K = self.k_proj(hidden_states).view(B, L, self.num_kv_heads, d).permute(0, 2, 1, 3)
        V = self.v_proj(hidden_states).view(B, L, self.num_kv_heads, d).permute(0, 2, 1, 3)
        is_cross = bool(kwargs.get("is_cross_attention", False))

        if position_embeddings is not None:
            cos, sin = position_embeddings
            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)

        scale = 1.0 / math.sqrt(d)
        if (
            (2 * self.window_size + 1) >= L
            and attention_mask is None
            and not output_attentions
            and head_mask is None
        ):
            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 = getattr(self, "is_decoder", False) and not is_cross and L > 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, L, self.all_head_size)
                )
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                return (
                    self.out_proj(_sdpa_out)
                    if not getattr(self, "is_nested_encoder", False)
                    else _sdpa_out
                ), None
            except Exception:
                pass

        scores = (
            torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
        )  # [B, H, L, L]

        # Valid mask handling
        valid_mask_2d = get_token_valid_mask(attention_mask, B, L, device)
        key_valid_4d = valid_mask_2d.unsqueeze(1).unsqueeze(2)

        # Local sliding window boolean mask
        idx = torch.arange(L, device=device)
        if self.is_decoder:
            local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0) <= self.window_size) & (
                idx.unsqueeze(1) >= idx.unsqueeze(0)
            )
        else:
            local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0)).abs() <= self.window_size
        is_dec = bool(
            self.is_decoder
            or getattr(self, "is_causal", False)
            or getattr(self.config, "is_decoder", False)
            or getattr(self.config, "is_causal", False)
        )
        if is_dec:
            local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0) <= self.window_size) & (
                idx.unsqueeze(1) >= idx.unsqueeze(0)
            )
        else:
            local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0)).abs() <= self.window_size
        local_mask_4d = local_mask.unsqueeze(0).unsqueeze(0)

        if (2 * self.window_size + 1) >= L:
            total_mask = local_mask_4d & key_valid_4d
        else:
            non_local_scores = scores.masked_fill(local_mask_4d | (~key_valid_4d), -1e4)
            if is_dec:
                causal_tril = (
                    torch.tril(torch.ones(L, L, device=device, dtype=torch.bool))
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                non_local_scores = non_local_scores.masked_fill(~causal_tril, -1e4)
            k_val = min(self.top_k, max(1, L - (2 * self.window_size + 1)))
            _, topk_ind = torch.topk(non_local_scores, k=k_val, dim=-1)
            topk_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_ind, True)
            total_mask = (local_mask_4d | topk_mask) & key_valid_4d

        if is_dec:
            causal_tril = (
                torch.tril(torch.ones(L, L, device=device, dtype=torch.bool))
                .unsqueeze(0)
                .unsqueeze(0)
            )
            total_mask = total_mask & causal_tril

        scores = scores.masked_fill(~total_mask, -1e4)
        attn = F.softmax(scores, dim=-1).to(dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        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.view(head_mask.size(0), -1, 1, 1)
            else:
                hm = head_mask
            attn = attn * hm

        context = torch.matmul(attn, V)
        context = context * valid_mask_2d.unsqueeze(1).unsqueeze(-1).to(dtype)
        context = context.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)

        return output, (attn if output_attentions else None)

Every 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.

Methods

def copy_weights(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    if src is None:
        return

    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")
    )
    if cls_name in ["BertSelfAttention", "RobertaSelfAttention"] or (
        hasattr(src, "query") and not has_out
    ):
        self.is_nested_encoder = True
    if hasattr(src, "self") and hasattr(src, "output"):
        self.is_nested_encoder = True

    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
    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()
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or hasattr(c_attn_mod, "nf")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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()
                if (
                    (c_proj_mod.__class__.__name__ == "Conv1D")
                    or (base_o.__class__.__name__ == "Conv1D")
                    or hasattr(c_proj_mod, "nf")
                    or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                ):
                    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. 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
    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"):
                        w_data = mod.weight.data.clone()
                        if (
                            (mod.__class__.__name__ == "Conv1D")
                            or hasattr(mod, "nf")
                            or (
                                w_data.ndim == 2
                                and proj_key == "o"
                                and w_data.shape[0] < w_data.shape[1]
                            )
                        ):
                            w_data = w_data.t().contiguous()
                        extracted[proj_key] = (
                            w_data,
                            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(f"UniversalAdapter fallback: {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"NativeSparseAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}."
        )

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()
            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device, dtype=target_mod.weight.dtype
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

            w_norm = float(torch.norm(target_mod.weight.data).item())
            if w_norm < 1e-4:
                logger.warning(
                    f"NativeSparseAttention weight verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                )

            if self.freeze_base:
                target_mod.freeze_base_weights()

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,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
output_attentions: bool | None = False,
*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,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    output_attentions: bool | None = False,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, L, D = hidden_states.shape
    H, d = self.num_heads, self.head_dim
    device = hidden_states.device
    dtype = hidden_states.dtype

    Q = self.q_proj(hidden_states).view(B, L, H, d).permute(0, 2, 1, 3)
    K = self.k_proj(hidden_states).view(B, L, self.num_kv_heads, d).permute(0, 2, 1, 3)
    V = self.v_proj(hidden_states).view(B, L, self.num_kv_heads, d).permute(0, 2, 1, 3)
    is_cross = bool(kwargs.get("is_cross_attention", False))

    if position_embeddings is not None:
        cos, sin = position_embeddings
        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)

    scale = 1.0 / math.sqrt(d)
    if (
        (2 * self.window_size + 1) >= L
        and attention_mask is None
        and not output_attentions
        and head_mask is None
    ):
        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 = getattr(self, "is_decoder", False) and not is_cross and L > 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, L, self.all_head_size)
            )
            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None
        except Exception:
            pass

    scores = (
        torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
    )  # [B, H, L, L]

    # Valid mask handling
    valid_mask_2d = get_token_valid_mask(attention_mask, B, L, device)
    key_valid_4d = valid_mask_2d.unsqueeze(1).unsqueeze(2)

    # Local sliding window boolean mask
    idx = torch.arange(L, device=device)
    if self.is_decoder:
        local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0) <= self.window_size) & (
            idx.unsqueeze(1) >= idx.unsqueeze(0)
        )
    else:
        local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0)).abs() <= self.window_size
    is_dec = bool(
        self.is_decoder
        or getattr(self, "is_causal", False)
        or getattr(self.config, "is_decoder", False)
        or getattr(self.config, "is_causal", False)
    )
    if is_dec:
        local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0) <= self.window_size) & (
            idx.unsqueeze(1) >= idx.unsqueeze(0)
        )
    else:
        local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0)).abs() <= self.window_size
    local_mask_4d = local_mask.unsqueeze(0).unsqueeze(0)

    if (2 * self.window_size + 1) >= L:
        total_mask = local_mask_4d & key_valid_4d
    else:
        non_local_scores = scores.masked_fill(local_mask_4d | (~key_valid_4d), -1e4)
        if is_dec:
            causal_tril = (
                torch.tril(torch.ones(L, L, device=device, dtype=torch.bool))
                .unsqueeze(0)
                .unsqueeze(0)
            )
            non_local_scores = non_local_scores.masked_fill(~causal_tril, -1e4)
        k_val = min(self.top_k, max(1, L - (2 * self.window_size + 1)))
        _, topk_ind = torch.topk(non_local_scores, k=k_val, dim=-1)
        topk_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_ind, True)
        total_mask = (local_mask_4d | topk_mask) & key_valid_4d

    if is_dec:
        causal_tril = (
            torch.tril(torch.ones(L, L, device=device, dtype=torch.bool))
            .unsqueeze(0)
            .unsqueeze(0)
        )
        total_mask = total_mask & causal_tril

    scores = scores.masked_fill(~total_mask, -1e4)
    attn = F.softmax(scores, dim=-1).to(dtype)
    attn = torch.nan_to_num(attn, nan=0.0)
    attn = self.dropout(attn)

    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.view(head_mask.size(0), -1, 1, 1)
        else:
            hm = head_mask
        attn = attn * hm

    context = torch.matmul(attn, V)
    context = context * valid_mask_2d.unsqueeze(1).unsqueeze(-1).to(dtype)
    context = context.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)

    if self.is_nested_encoder:
        output = context
    else:
        output = self.out_proj(context)

    return output, (attn if output_attentions else None)

Must return (output, attention_weights_or_None).

class NativeSparseAttention (config: Any | None = None,
rank: int | None = None,
window_size: int = 32,
top_k: int = 16,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "nsa", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "nsa_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class NativeSparseAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        window_size: int = 32,
        top_k: int = 16,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        hidden_size = kwargs.get("hidden_size")
        num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")
        num_kv_heads = kwargs.get("num_key_value_heads") or kwargs.get("num_kv_heads")

        if hidden_size is None:
            if config is not None:
                hidden_size = getattr(config, "hidden_size", getattr(config, "dim", 768))
            elif source_attn is not None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                hidden_size = (
                    getattr(out_layer, "out_features", 768) if out_layer is not None else 768
                )
            else:
                hidden_size = 768

        if num_heads is None:
            if config is not None:
                num_heads = getattr(config, "num_attention_heads", getattr(config, "n_heads", 12))
            elif source_attn is not None:
                num_heads = getattr(
                    source_attn,
                    "num_attention_heads",
                    getattr(source_attn, "n_heads", 12),
                )
            else:
                num_heads = 12

        if num_kv_heads is None:
            if config is not None:
                num_kv_heads = getattr(
                    config,
                    "num_key_value_heads",
                    getattr(config, "num_kv_heads", num_heads),
                )
            elif source_attn is not None:
                num_kv_heads = getattr(source_attn, "num_key_value_heads", num_heads)
            else:
                num_kv_heads = 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 = self.dim // self.num_heads if self.num_heads > 0 else 64
        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.window_size = (
            getattr(config, "nsa_window_size", window_size) if config else window_size
        )
        self.top_k = getattr(config, "nsa_top_k", top_k) if config else top_k

        # Decoder detection
        model_type = getattr(config, "model_type", "").lower() if config else ""
        is_dec = (
            getattr(
                config,
                "is_decoder",
                getattr(config, "is_causal", kwargs.get("is_causal", False)),
            )
            if config
            else False
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(_mt.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        # LoRA config
        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0

        # Projections
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        if source_attn is not None:
            self.copy_weights(source_attn)

    def copy_weights(self, src: Any) -> None:
        if src is None:
            return

        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")
        )
        if cls_name in ["BertSelfAttention", "RobertaSelfAttention"] or (
            hasattr(src, "query") and not has_out
        ):
            self.is_nested_encoder = True
        if hasattr(src, "self") and hasattr(src, "output"):
            self.is_nested_encoder = True

        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
        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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or hasattr(c_attn_mod, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (c_proj_mod.__class__.__name__ == "Conv1D")
                        or (base_o.__class__.__name__ == "Conv1D")
                        or hasattr(c_proj_mod, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        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. 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
        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"):
                            w_data = mod.weight.data.clone()
                            if (
                                (mod.__class__.__name__ == "Conv1D")
                                or hasattr(mod, "nf")
                                or (
                                    w_data.ndim == 2
                                    and proj_key == "o"
                                    and w_data.shape[0] < w_data.shape[1]
                                )
                            ):
                                w_data = w_data.t().contiguous()
                            extracted[proj_key] = (
                                w_data,
                                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(f"UniversalAdapter fallback: {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"NativeSparseAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                w_norm = float(torch.norm(target_mod.weight.data).item())
                if w_norm < 1e-4:
                    logger.warning(
                        f"NativeSparseAttention weight verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                    )

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    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,
        output_attentions: bool | None = False,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = self.q_proj(hidden_states).view(B, L, H, d).permute(0, 2, 1, 3)
        K = self.k_proj(hidden_states).view(B, L, self.num_kv_heads, d).permute(0, 2, 1, 3)
        V = self.v_proj(hidden_states).view(B, L, self.num_kv_heads, d).permute(0, 2, 1, 3)
        is_cross = bool(kwargs.get("is_cross_attention", False))

        if position_embeddings is not None:
            cos, sin = position_embeddings
            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)

        scale = 1.0 / math.sqrt(d)
        if (
            (2 * self.window_size + 1) >= L
            and attention_mask is None
            and not output_attentions
            and head_mask is None
        ):
            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 = getattr(self, "is_decoder", False) and not is_cross and L > 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, L, self.all_head_size)
                )
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                return (
                    self.out_proj(_sdpa_out)
                    if not getattr(self, "is_nested_encoder", False)
                    else _sdpa_out
                ), None
            except Exception:
                pass

        scores = (
            torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
        )  # [B, H, L, L]

        # Valid mask handling
        valid_mask_2d = get_token_valid_mask(attention_mask, B, L, device)
        key_valid_4d = valid_mask_2d.unsqueeze(1).unsqueeze(2)

        # Local sliding window boolean mask
        idx = torch.arange(L, device=device)
        if self.is_decoder:
            local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0) <= self.window_size) & (
                idx.unsqueeze(1) >= idx.unsqueeze(0)
            )
        else:
            local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0)).abs() <= self.window_size
        is_dec = bool(
            self.is_decoder
            or getattr(self, "is_causal", False)
            or getattr(self.config, "is_decoder", False)
            or getattr(self.config, "is_causal", False)
        )
        if is_dec:
            local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0) <= self.window_size) & (
                idx.unsqueeze(1) >= idx.unsqueeze(0)
            )
        else:
            local_mask = (idx.unsqueeze(1) - idx.unsqueeze(0)).abs() <= self.window_size
        local_mask_4d = local_mask.unsqueeze(0).unsqueeze(0)

        if (2 * self.window_size + 1) >= L:
            total_mask = local_mask_4d & key_valid_4d
        else:
            non_local_scores = scores.masked_fill(local_mask_4d | (~key_valid_4d), -1e4)
            if is_dec:
                causal_tril = (
                    torch.tril(torch.ones(L, L, device=device, dtype=torch.bool))
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                non_local_scores = non_local_scores.masked_fill(~causal_tril, -1e4)
            k_val = min(self.top_k, max(1, L - (2 * self.window_size + 1)))
            _, topk_ind = torch.topk(non_local_scores, k=k_val, dim=-1)
            topk_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_ind, True)
            total_mask = (local_mask_4d | topk_mask) & key_valid_4d

        if is_dec:
            causal_tril = (
                torch.tril(torch.ones(L, L, device=device, dtype=torch.bool))
                .unsqueeze(0)
                .unsqueeze(0)
            )
            total_mask = total_mask & causal_tril

        scores = scores.masked_fill(~total_mask, -1e4)
        attn = F.softmax(scores, dim=-1).to(dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        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.view(head_mask.size(0), -1, 1, 1)
            else:
                hm = head_mask
            attn = attn * hm

        context = torch.matmul(attn, V)
        context = context * valid_mask_2d.unsqueeze(1).unsqueeze(-1).to(dtype)
        context = context.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)

        return output, (attn if output_attentions else None)

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class NystromAttention (config: Any | None = None,
rank: int | None = None,
num_landmarks: int | None = None,
pinv_iterations: int = 3,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "nystrom", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "nystrom_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class NystromAttention(BaseAttention):
    is_custom_attention: bool = True
    triton_kernel: str = "linear_attn"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        num_landmarks: int | None = None,
        pinv_iterations: int = 3,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        hidden_size = kwargs.get("hidden_size")
        num_heads = kwargs.get("num_heads") or kwargs.get("num_attention_heads")
        num_kv_heads = kwargs.get("num_key_value_heads") or kwargs.get("num_kv_heads")

        if hidden_size is None:
            if config is not None:
                hidden_size = getattr(config, "hidden_size", getattr(config, "dim", 768))
            elif source_attn is not None:
                out_layer = getattr(
                    source_attn,
                    "out_lin",
                    getattr(
                        source_attn,
                        "dense",
                        getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)),
                    ),
                )
                hidden_size = (
                    getattr(out_layer, "out_features", 768) if out_layer is not None else 768
                )
            else:
                hidden_size = 768

        if num_heads is None:
            if config is not None:
                num_heads = getattr(config, "num_attention_heads", getattr(config, "n_heads", 12))
            elif source_attn is not None:
                num_heads = getattr(
                    source_attn,
                    "num_attention_heads",
                    getattr(source_attn, "n_heads", 12),
                )
            else:
                num_heads = 12

        if num_kv_heads is None:
            if config is not None:
                num_kv_heads = getattr(
                    config,
                    "num_key_value_heads",
                    getattr(config, "num_kv_heads", num_heads),
                )
            elif source_attn is not None:
                num_kv_heads = getattr(source_attn, "num_key_value_heads", num_heads)
            else:
                num_kv_heads = 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 = self.dim // self.num_heads if self.num_heads > 0 else 64
        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

        landmarks_cfg = (
            (
                num_landmarks
                if num_landmarks is not None
                else getattr(
                    config,
                    "num_landmarks",
                    getattr(config, "nystrom_landmarks", rank if rank is not None else 32),
                )
            )
            if config is not None
            else (num_landmarks if num_landmarks is not None else 32)
        )
        self.num_landmarks = min(32, int(landmarks_cfg))
        self.pinv_iterations = pinv_iterations

        # Decoder detection
        model_type = getattr(config, "model_type", "").lower() if config else ""
        is_dec = (
            getattr(
                config,
                "is_decoder",
                getattr(config, "is_causal", kwargs.get("is_causal", False)),
            )
            if config
            else False
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(_mt.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        # LoRA config
        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0

        # Projections
        self.query = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.key = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.value = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)
        self.triton_kwargs = {"phi": "identity", "eps": 1e-6}

        if source_attn is not None:
            self.copy_weights(source_attn)

    @property
    def q_proj(self):
        return self.query

    @property
    def k_proj(self):
        return self.key

    @property
    def v_proj(self):
        return self.value

    @property
    def q_lin(self):
        return self.query

    @property
    def k_lin(self):
        return self.key

    @property
    def v_lin(self):
        return self.value

    @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:
        if src is None:
            return

        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)

        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
        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_conv1d = (
                    "Conv1D" in c_attn_mod.__class__.__name__
                    or "Conv1D" in base_c.__class__.__name__
                    or hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                )
                if is_conv1d:
                    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_conv1d = (
                        "Conv1D" in c_proj_mod.__class__.__name__
                        or "Conv1D" in base_o.__class__.__name__
                        or hasattr(c_proj_mod, "nf")
                        or hasattr(base_o, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    )
                    if is_conv1d:
                        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. 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
        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:
                pass

        target_map = {
            "q": self.query,
            "k": self.key,
            "v": self.value,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()
                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                w_norm = float(torch.linalg.norm(target_mod.weight.data.float()).item())
                if w_norm < 1e-4:
                    logger.warning(
                        f"NystromAttention weight verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                    )

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    def _iterative_pinv(self, A: torch.Tensor, iters: int = 3) -> torch.Tensor:
        try:
            return torch.linalg.pinv(A.float()).to(A.dtype)
        except Exception:
            norm_1 = A.abs().sum(dim=-2).amax(dim=-1, keepdim=True).unsqueeze(-1)
            norm_inf = A.abs().sum(dim=-1).amax(dim=-1, keepdim=True).unsqueeze(-1)
            alpha = 1.0 / (norm_1 * norm_inf + 1e-5)

            V = alpha * A.transpose(-2, -1).contiguous()
            I = torch.eye(A.size(-1), device=A.device, dtype=A.dtype).unsqueeze(0).unsqueeze(0)

            for _ in range(iters):
                V = torch.matmul(V, 2.0 * I - torch.matmul(A, V))
            return V

    def _dense_attention(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        valid_mask: torch.Tensor,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        scale = 1.0 / math.sqrt(self.head_dim)
        try:
            if valid_mask.all():
                return F.scaled_dot_product_attention(
                    Q,
                    K,
                    V,
                    is_causal=getattr(self, "is_causal", False)
                    or getattr(self, "is_decoder", False),
                    dropout_p=self.dropout.p if self.training else 0.0,
                ).to(dtype)
            else:
                attn_mask = valid_mask.unsqueeze(1).unsqueeze(2)
                return F.scaled_dot_product_attention(
                    Q,
                    K,
                    V,
                    attn_mask=attn_mask,
                    dropout_p=self.dropout.p if self.training else 0.0,
                ).to(dtype)
        except Exception:
            pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
        scores = scores.masked_fill(~valid_mask.unsqueeze(1).unsqueeze(2), -1e4)
        attn_weights = F.softmax(scores, dim=-1).to(dtype)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights)
        return torch.matmul(attn_weights, V)

    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,
        output_attentions: bool | None = False,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if kwargs.get("encoder_hidden_states") is not None:
            raise NotImplementedError("Cross-attention is not implemented for NystromAttention")

        B, N, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        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 not None:
            cos, sin = position_embeddings
            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)

        valid_mask = get_token_valid_mask(attention_mask, B, N, device)

        if self.is_decoder or self.is_causal:
            scale = 1.0 / math.sqrt(d)
            scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
            causal_mask = torch.tril(torch.ones((N, N), device=device, dtype=torch.bool))
            scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4)
            scores = scores.masked_fill(~valid_mask.unsqueeze(1).unsqueeze(2), -1e4)
            attn_weights = F.softmax(scores, dim=-1).to(dtype)
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            attn_weights = self.dropout(attn_weights)
            context = torch.matmul(attn_weights, V)
        elif max(128, 2 * self.num_landmarks) >= N:
            # Short sequence dense fallback: Nyström low-rank approximation is mathematically
            # sub-optimal and memory-inefficient when N <= 2*landmarks (or N <= 128)
            context = self._dense_attention(Q, K, V, valid_mask, dtype)
        else:
            m = min(self.num_landmarks, N)

            # Adaptive pooling landmarks
            Q_flat = Q.reshape(B * H, N, d).transpose(1, 2).contiguous()
            K_flat = K.reshape(B * H, N, d).transpose(1, 2).contiguous()

            Q_landmarks = (
                F.adaptive_avg_pool1d(Q_flat, m).transpose(1, 2).contiguous().reshape(B, H, m, d)
            )
            K_landmarks = (
                F.adaptive_avg_pool1d(K_flat, m).transpose(1, 2).contiguous().reshape(B, H, m, d)
            )

            scale = 1.0 / math.sqrt(d)
            scores1 = (
                torch.matmul(Q.float(), K_landmarks.float().transpose(-2, -1).contiguous()) * scale
            )
            scores2 = (
                torch.matmul(
                    Q_landmarks.float(),
                    K_landmarks.float().transpose(-2, -1).contiguous(),
                )
                * scale
            )
            scores3 = (
                torch.matmul(Q_landmarks.float(), K.float().transpose(-2, -1).contiguous()) * scale
            )

            mask_expanded = valid_mask.unsqueeze(1).unsqueeze(2)
            scores3 = scores3.masked_fill(~mask_expanded, -1e4)

            F1 = F.softmax(scores1, dim=-1).to(dtype)
            F2 = F.softmax(scores2, dim=-1).to(dtype)
            F3 = F.softmax(scores3, dim=-1).to(dtype)

            F2_pinv = self._iterative_pinv(F2, iters=self.pinv_iterations)

            # Associative linear multiplication: O(N * m * d)
            step1 = torch.matmul(F3, V)
            step2 = torch.matmul(F2_pinv, step1)
            context = torch.matmul(F1, step2)

            context = torch.nan_to_num(context, nan=0.0)
            context = self.dropout(context)

        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.view(head_mask.size(0), -1, 1, 1)
            else:
                hm = head_mask
            context = context * hm

        context = context * valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_proj(context)

        return out, None

Every 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 triton_kernel : str

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.key

prop k_proj

Expand source code

@property
def k_proj(self):
    return self.key

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.query

prop q_proj

Expand source code

@property
def q_proj(self):
    return self.query

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.value

prop v_proj

Expand source code

@property
def v_proj(self):
    return self.value

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class ORIONAttention (config: Any | None = None,
rank: int | None = None,
model_config: Any | None = None,
orion_cfg: ORIONConfig | None = None,
pretrained_attn: Any | None = None,
old_attn: Any | None = None,
comp_config: Any | None = None,
bert_mode: bool = False,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("orion_module", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("orion", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("orion_attention", supports_causal=True, supports_gqa=True)
class ORIONAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        model_config: Any | None = None,
        orion_cfg: ORIONConfig | None = None,
        pretrained_attn: Any | None = None,
        old_attn: Any | None = None,
        comp_config: Any | None = None,
        bert_mode: bool = False,
        **kwargs: Any,
    ) -> None:
        cfg = config or model_config
        source_attn = (
            pretrained_attn
            if pretrained_attn is not None
            else (old_attn if old_attn is not None else kwargs.get("source_attn"))
        )
        try:
            super().__init__(cfg, rank)
        except TypeError:
            super().__init__(cfg)

        if orion_cfg is None:
            orion_cfg = ORIONConfig()
            if rank is not None:
                orion_cfg.svd_max_rank = max(rank, orion_cfg.svd_min_rank)

        hidden_size = kwargs.get("hidden_size")
        num_heads = kwargs.get("num_attention_heads", kwargs.get("num_heads"))
        num_kv_heads = kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads"))

        if hidden_size is None or num_heads is None:
            if cfg is not None:
                if hidden_size is None:
                    hidden_size = _cfg_get(
                        cfg,
                        "hidden_size",
                        _cfg_get(cfg, "d_model", _cfg_get(cfg, "dim", 768)),
                    )
                if num_heads is None:
                    num_heads = _cfg_get(cfg, "num_attention_heads", _cfg_get(cfg, "num_heads", 12))
                if num_kv_heads is None:
                    num_kv_heads = _cfg_get(
                        cfg,
                        "num_key_value_heads",
                        _cfg_get(cfg, "num_kv_heads", num_heads),
                    )
            elif source_attn is not None:
                if hidden_size is None:
                    out_layer = getattr(
                        source_attn,
                        "out_lin",
                        getattr(
                            source_attn,
                            "dense",
                            getattr(
                                source_attn,
                                "c_proj",
                                getattr(source_attn, "o_proj", None),
                            ),
                        ),
                    )
                    hidden_size = (
                        getattr(out_layer, "out_features", 768) if out_layer is not None else 768
                    )
                if num_heads is None:
                    num_heads = getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(source_attn, "n_heads", 12),
                    )
                if num_kv_heads is None:
                    num_kv_heads = getattr(source_attn, "num_key_value_heads", num_heads)
            else:
                hidden_size = 768 if hidden_size is None else hidden_size
                num_heads = 12 if num_heads is None else num_heads
                num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads

        if num_kv_heads is None:
            num_kv_heads = num_heads

        self.num_heads = max(1, int(num_heads))
        self.num_kv_heads = max(1, int(num_kv_heads))
        self.hidden_size = int(hidden_size)
        if self.hidden_size % self.num_heads != 0:
            raise ValueError(
                f"hidden_size {self.hidden_size} must be divisible by num_heads {self.num_heads}"
            )
        self.head_dim = self.hidden_size // self.num_heads
        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.cfg = orion_cfg
        self.bert_mode = bert_mode

        # Allocate kernel vs full heads
        if self.num_heads <= 1:
            self.num_kernel_heads = 1
            self.num_full_heads = 0
        else:
            self.num_full_heads = max(1, self.num_heads // 3)
            self.num_kernel_heads = self.num_heads - self.num_full_heads

        # Decoder detection
        model_type = str(getattr(cfg, "model_type", "") or "").lower() if cfg else ""
        is_dec = (
            bool(
                getattr(
                    cfg,
                    "is_decoder",
                    getattr(cfg, "is_causal", kwargs.get("is_causal", False)),
                )
            )
            if cfg
            else False
        )
        if (
            any(
                model_type.startswith(p)
                for p in ("gpt", "qwen", "llama", "mistral", "falcon", "bloom", "opt")
            )
            or hasattr(cfg, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(cfg, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        # LoRA config
        self.lora_rank = (
            getattr(
                cfg,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if cfg
            else 16
        )
        self.lora_alpha = (
            getattr(
                cfg,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if cfg
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and source_attn is not None:
            p = next(source_attn.parameters(), None)
            if p is not None:
                target_dtype = p.dtype

        # Projections
        self.q_proj = FactorisedLinear(
            hidden_size,
            self.all_head_size,
            bias=True,
            energy=orion_cfg.svd_energy,
            min_rank=orion_cfg.svd_min_rank,
            max_rank=orion_cfg.svd_max_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.k_proj = FactorisedLinear(
            hidden_size,
            self.all_kv_head_size,
            bias=True,
            energy=orion_cfg.svd_energy,
            min_rank=orion_cfg.svd_min_rank,
            max_rank=orion_cfg.svd_max_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.v_proj = FactorisedLinear(
            hidden_size,
            self.all_kv_head_size,
            bias=True,
            energy=orion_cfg.svd_energy,
            min_rank=orion_cfg.svd_min_rank,
            max_rank=orion_cfg.svd_max_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.out_proj = FactorisedLinear(
            self.all_head_size,
            hidden_size,
            bias=True,
            energy=orion_cfg.svd_energy,
            min_rank=orion_cfg.svd_min_rank,
            max_rank=orion_cfg.svd_max_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )

        rff_init = torch.randn(
            max(1, self.num_kernel_heads), self.head_dim, orion_cfg.rff_features
        ) * (1.0 / math.sqrt(self.head_dim))
        self.register_buffer("rff_proj", rff_init, persistent=False)
        self.dropout = nn.Dropout(orion_cfg.dropout)

        if source_attn is not None:
            self.copy_weights(source_attn)

        if target_dtype is not None:
            self.to(target_dtype)

    # ------------------------------------------------------------------
    @property
    def query(self) -> FactorisedLinear:
        return self.q_proj

    @property
    def key(self) -> FactorisedLinear:
        return self.k_proj

    @property
    def value(self) -> FactorisedLinear:
        return self.v_proj

    @property
    def dense(self) -> FactorisedLinear:
        return self.out_proj

    @property
    def o_proj(self) -> FactorisedLinear:
        return self.out_proj

    @property
    def q_lin(self) -> FactorisedLinear:
        return self.q_proj

    @property
    def k_lin(self) -> FactorisedLinear:
        return self.k_proj

    @property
    def v_lin(self) -> FactorisedLinear:
        return self.v_proj

    @property
    def out_lin(self) -> FactorisedLinear:
        return self.out_proj

    @property
    def c_proj(self) -> FactorisedLinear:
        return self.out_proj

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    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")
            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),
        }

        # GPT-2 fused
        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()
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        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])
                    ):
                        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,
                    )

        # 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,
                )

        # Flat modules (DistilBERT, Qwen, etc.)
        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

        # UniversalAdapter fallback
        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:
            logger.warning(
                "ORIONAttention.copy_weights: missing projections %s from %s",
                missing_after,
                src.__class__.__name__,
            )
            return

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        # Identity for nested encoders
                        target_mod.down.weight.data.zero_()
                        target_mod.up.weight.data.zero_()
                        r = min(target_mod.rank, self.hidden_size, self.all_head_size)
                        target_mod.down.weight.data[:r, :r].fill_diagonal_(1.0)
                        target_mod.up.weight.data[:r, :r].fill_diagonal_(1.0)
                        if target_mod.up.bias is not None:
                            target_mod.up.bias.data.zero_()
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                # Orientation
                expected = (target_mod.out_features, target_mod.in_features)
                if w.shape != expected and w.t().contiguous().shape == expected:
                    w = w.t().contiguous()

                target_mod.decompose_from_weight(
                    w[: target_mod.out_features, : target_mod.in_features],
                    b[: target_mod.out_features] if b is not None else None,
                )

                norm = float(torch.linalg.norm(target_mod.weight.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "ORIONAttention: %s has near-zero Frobenius norm (%.2e)",
                        proj_key,
                        norm,
                    )

        # ------------------------------------------------------------------
        # OR-FIX-03: _kernel_linear_attention with token_valid
        # ------------------------------------------------------------------

    def _kernel_linear_attention(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        token_valid: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        FAVOR+ random feature linear attention.

        Implements the unbiased feature map from Choromanski et al. (2020):
            φ(x) = exp(ω·x - ||x||²/2) / sqrt(m),   ω ~ N(0, I_d)

        The `-||x||²/2` term is essential for an unbiased estimate of exp(q·k).
        The scaling 1/sqrt(d) is applied to x once (not to ω, not twice).
        """
        B, Hk, N, D = Q.shape
        if Hk == 0:
            return torch.zeros(B, 0, N, D, device=Q.device, dtype=Q.dtype)

        F_dim = self.cfg.rff_features
        orig_dtype = Q.dtype
        Q, K, V = Q.float(), K.float(), V.float()
        proj = self.rff_proj[:Hk].to(Q.device).float()  # (Hk, D, m); ω ~ N(0, I_d)

        # Scale x ONCE by 1/sqrt(d). This is the "x" entering the FAVOR+ formula.
        scale = 1.0 / math.sqrt(self.head_dim)
        Q_scaled = Q * scale
        K_scaled = K * scale

        # Project onto random features:  (B, Hk, N, m)
        Q_proj = torch.einsum("bhnd,hdf->bhnf", Q_scaled, proj)
        K_proj = torch.einsum("bhnd,hdf->bhnf", K_scaled, proj)

        # Subtract ||x||²/2 (per query/key) → unbiased FAVOR+ feature map
        Q_norm_sq = (Q_scaled**2).sum(dim=-1, keepdim=True) / 2.0  # (B, Hk, N, 1)
        K_norm_sq = (K_scaled**2).sum(dim=-1, keepdim=True) / 2.0  # (B, Hk, N, 1)

        # Numerical safety: subtract global max before exp (log-sum-exp style)
        # Mathematically cancels out in num/denom, so accuracy is unchanged.
        max_val = max(float(Q_proj.detach().amax().item()), float(K_proj.detach().amax().item()))
        max_val = max(max_val, 0.0)  # only subtract if positive

        Q_feat = torch.exp(Q_proj - Q_norm_sq - max_val) / math.sqrt(F_dim)
        K_feat = torch.exp(K_proj - K_norm_sq - max_val) / math.sqrt(F_dim)

        # Padding mask on K and V
        if token_valid is not None:
            mask_kv = token_valid.unsqueeze(1).unsqueeze(-1).to(K_feat.dtype)
            K_feat = K_feat * mask_kv
            V = V * mask_kv

        M = Q_feat.shape[-1]
        if self.is_decoder:
            chunk_size = min(64, N)
            num_chunks = (N + chunk_size - 1) // chunk_size
            S = torch.zeros(B, Hk, M, D, dtype=torch.float32, device=Q.device)
            Z = torch.zeros(B, Hk, M, dtype=torch.float32, device=Q.device)
            ctx_list = []

            for c in range(num_chunks):
                start = c * chunk_size
                end = min(start + chunk_size, N)
                L = end - start

                q_c = Q_feat[:, :, start:end]
                k_c = K_feat[:, :, start:end]
                v_c = V[:, :, start:end]

                causal_mask = torch.tril(torch.ones(L, L, device=Q.device, dtype=torch.float32))
                intra_attn = torch.matmul(q_c, k_c.transpose(-1, -2).contiguous()) * causal_mask
                intra_out = torch.matmul(intra_attn, v_c)
                inter_out = torch.matmul(q_c, S)

                intra_den = intra_attn.sum(dim=-1, keepdim=True)
                inter_den = torch.matmul(q_c, Z.unsqueeze(-1))

                den = (intra_den + inter_den).clamp(min=1e-4)
                ctx_list.append((intra_out + inter_out) / den)

                S = S + torch.matmul(k_c.transpose(-1, -2).contiguous(), v_c)
                Z = Z + k_c.sum(dim=2)

            ctx = torch.cat(ctx_list, dim=2)
        else:
            KV = torch.einsum("bhnf,bhnd->bhfd", K_feat, V)
            K_sum = K_feat.sum(dim=2)
            denom = torch.einsum("bhnf,bhf->bhn", Q_feat, K_sum).unsqueeze(-1).clamp(min=1e-4)
            ctx = torch.einsum("bhnf,bhfd->bhnd", Q_feat, KV) / denom
        return ctx.to(orig_dtype)

    # ------------------------------------------------------------------
    # OR-FIX-02: _full_softpick_attention rewritten
    # ------------------------------------------------------------------
    def _full_softpick_attention(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        token_valid: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        Full softmax attention for the "full" heads.

        Returns a tensor of shape (B, Hf, N, D). No internal out_proj is
        applied here; the caller handles out_proj after concatenating all
        head groups.
        """
        B, Hf, N, D = Q.shape
        if Hf == 0:
            return torch.zeros(B, 0, N, D, device=Q.device, dtype=Q.dtype)

        scale = self.head_dim**-0.5
        is_causal = bool(self.is_decoder) and N > 1

        # Build additive attention bias for padding
        attn_bias = None
        if token_valid is not None:
            valid_bias = ((1.0 - token_valid.float()) * -1e4).view(B, 1, 1, N)
            attn_bias = valid_bias.to(Q.dtype)

        # Try 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
            bias_sdp = None
            if attn_bias is not None:
                bias_sdp = attn_bias.to(q_sdp.dtype)

            out = F.scaled_dot_product_attention(
                q_sdp,
                k_sdp,
                v_sdp,
                attn_mask=bias_sdp,
                dropout_p=self.dropout.p if self.training else 0.0,
                is_causal=is_causal and bias_sdp is None,
            )
            out = out.to(Q.dtype)
            return torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
        except Exception:
            pass

        # Fallback: manual temperature softmax
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if is_causal:
            causal_mask = torch.tril(torch.ones(N, N, device=Q.device, dtype=torch.bool))
            scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4)

        if token_valid is not None:
            scores = scores.masked_fill(~token_valid.unsqueeze(1).unsqueeze(2).bool(), -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_w = temperature_softmax(scores, dim=-1, temperature=self.cfg.softpick_temperature).to(
            V.dtype
        )
        attn_w = self.dropout(attn_w)
        return torch.matmul(attn_w, V)

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if hidden_states.dim() == 4:
            B, H, L, Dh = hidden_states.shape
            hidden_states = hidden_states.view(B, L, H * Dh)
        elif not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        B, N, _ = hidden_states.shape
        Hk = self.num_kernel_heads
        Hf = self.num_full_heads
        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 = kv_input.shape[1]

        Q_all = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K_all = (
            self.k_proj(kv_input)
            .view(B, N_kv, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V_all = (
            self.v_proj(kv_input)
            .view(B, N_kv, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # RoPE
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            Q_all, K_all = apply_rotary_pos_emb(Q_all, K_all, cos, sin)

        # GQA expansion
        if self.kv_repeat > 1:
            K_all = repeat_kv(K_all, self.kv_repeat)
            V_all = repeat_kv(V_all, self.kv_repeat)

        # Token valid mask
        token_valid = None
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            token_valid = get_token_valid_mask(eff_mask, B, N_kv, device)

        ctx_list = []

        # OR-FIX-04: pass token_valid to both branches
        if Hk > 0:
            Q_kernel = Q_all[:, :Hk]
            K_kernel = K_all[:, :Hk]
            V_kernel = V_all[:, :Hk]
            ctx_kernel = self._kernel_linear_attention(Q_kernel, K_kernel, V_kernel, token_valid)
            ctx_list.append(ctx_kernel)

        if Hf > 0:
            Q_full = Q_all[:, Hk:]
            K_full = K_all[:, Hk:]
            V_full = V_all[:, Hk:]
            ctx_full = self._full_softpick_attention(Q_full, K_full, V_full, token_valid)
            ctx_list.append(ctx_full)

        if len(ctx_list) == 1:
            context = ctx_list[0]
        else:
            context = torch.cat(ctx_list, dim=1)

        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # 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.view(head_mask.size(0), -1, 1, 1)
            else:
                hm = head_mask
            context = context * hm.to(context.dtype)

        # Reshape to (B, N, H*D)
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # Output projection (skipped for nested encoders)
        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)
        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)

        return output, None

    def get_ortho_loss(self) -> torch.Tensor:
        return torch.zeros((), device=self.out_proj.up.weight.device)

Every 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.

Instance variables

prop c_proj : FactorisedLinear

Expand source code

@property
def c_proj(self) -> FactorisedLinear:
    return self.out_proj

prop dense : FactorisedLinear

Expand source code

@property
def dense(self) -> FactorisedLinear:
    return self.out_proj

prop k_lin : FactorisedLinear

Expand source code

@property
def k_lin(self) -> FactorisedLinear:
    return self.k_proj

prop key : FactorisedLinear

Expand source code

@property
def key(self) -> FactorisedLinear:
    return self.k_proj

prop o_proj : FactorisedLinear

Expand source code

@property
def o_proj(self) -> FactorisedLinear:
    return self.out_proj

prop out_lin : FactorisedLinear

Expand source code

@property
def out_lin(self) -> FactorisedLinear:
    return self.out_proj

prop q_lin : FactorisedLinear

Expand source code

@property
def q_lin(self) -> FactorisedLinear:
    return self.q_proj

prop query : FactorisedLinear

Expand source code

@property
def query(self) -> FactorisedLinear:
    return self.q_proj

prop v_lin : FactorisedLinear

Expand source code

@property
def v_lin(self) -> FactorisedLinear:
    return self.v_proj

prop value : FactorisedLinear

Expand source code

@property
def value(self) -> FactorisedLinear:
    return self.v_proj

Methods

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    return torch.zeros((), device=self.out_proj.up.weight.device)

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class ORIONConfig (svd_energy: float = 0.95,
svd_min_rank: int = 4,
svd_max_rank: int = 64,
num_kernel_heads: int = 8,
num_full_heads: int = 4,
rff_features: int = 256,
softpick_temperature: float = 1.0,
dropout: float = 0.1)

Expand source code

@dataclass
class ORIONConfig:
    svd_energy: float = 0.95
    svd_min_rank: int = 4
    svd_max_rank: int = 64
    num_kernel_heads: int = 8
    num_full_heads: int = 4
    rff_features: int = 256
    softpick_temperature: float = 1.0
    dropout: float = 0.1

ORIONConfig(svd_energy: 'float' = 0.95, svd_min_rank: 'int' = 4, svd_max_rank: 'int' = 64, num_kernel_heads: 'int' = 8, num_full_heads: 'int' = 4, rff_features: 'int' = 256, softpick_temperature: 'float' = 1.0, dropout: 'float' = 0.1)

Instance variables

var dropout : float

The type of the None singleton.

var num_full_heads : int

The type of the None singleton.

var num_kernel_heads : int

The type of the None singleton.

var rff_features : int

The type of the None singleton.

var softpick_temperature : float

The type of the None singleton.

var svd_energy : float

The type of the None singleton.

var svd_max_rank : int

The type of the None singleton.

var svd_min_rank : int

The type of the None singleton.

class OptimizedHybridAttention (config=None,
rank=None,
window_size=64,
compression_rank=32,
use_gate=True,
gate_bias=2.0,
pretrained_attn=None,
old_attn=None,
comp_config=None,
is_decoder=None,
**kwargs)

Expand source code

@ATTENTION_REGISTRY.register("oha", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("optimized_hybrid_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("oha_attention", supports_causal=True, supports_gqa=True)
class OptimizedHybridAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config=None,
        rank=None,
        window_size=64,
        compression_rank=32,
        use_gate=True,
        gate_bias=2.0,
        pretrained_attn=None,
        old_attn=None,
        comp_config=None,
        is_decoder=None,
        **kwargs,
    ):
        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 = (
            kwargs.get("head_dim")
            or getattr(config, "head_dim", None)
            or (hidden_size // max(1, 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.window_size = int(
            getattr(config, "oha_window_size", window_size) if config else window_size
        )
        self.window_radius = max(1, self.window_size // 2)

        cfg_cr = getattr(comp_config, "compression_rank", None) if comp_config else None
        self.compression_rank = int(cfg_cr if cfg_cr is not None else compression_rank)
        self.use_gate = bool(use_gate)

        self.is_decoder = _detect_decoder(
            config, is_decoder or kwargs.get("is_causal"), 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.freeze_base = self.lora_rank > 0

        self.query = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.key = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.value = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.q_proj = self.query
        self.k_proj = self.key
        self.v_proj = self.value

        self.compress_k = nn.Linear(self.dim, self.compression_rank, bias=False)
        self.compress_v = nn.Linear(self.dim, self.compression_rank, bias=False)
        nn.init.normal_(self.compress_k.weight, std=0.02)
        nn.init.normal_(self.compress_v.weight, std=0.02)

        if self.use_gate:
            self.gate = nn.Linear(self.dim, self.num_heads)
            nn.init.constant_(self.gate.bias, gate_bias)
            nn.init.normal_(self.gate.weight, std=0.02)
        else:
            self.register_parameter("gate", None)

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        if source_attn is not None:
            self.copy_weights(source_attn)

    # ---- detection helpers ----
    @staticmethod
    def _detect_hidden_size(config, src, kwargs):
        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_lin", "q_proj", "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, src, kwargs, hidden_size):
        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, src, kwargs, num_heads):
        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

    # ---- copy_weights (OHA-FIX-15) ----
    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")
            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 = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 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[0] < c_w.shape[1])
                )
                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,
                    )

        # RoBERTa / BERT 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,
                )

        # Flat modules
        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"OHA.copy_weights failed: missing mandatory projections "
                f"{missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.query,
            "k": self.key,
            "v": self.value,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue
                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                # OHA-FIX-15: log transposition
                if (
                    w.shape != target_mod.weight.shape
                    and w.t().contiguous().shape == target_mod.weight.shape
                ):
                    logger.debug(
                        "OHA.copy_weights: transposing '%s' from %s to %s",
                        proj_key,
                        tuple(w.shape),
                        tuple(target_mod.weight.shape),
                    )
                    w = w.t().contiguous()

                # OHA-FIX-15: log truncation
                if w.shape != target_mod.weight.shape:
                    logger.warning(
                        "OHA.copy_weights: '%s' shape mismatch %s vs target %s; truncating.",
                        proj_key,
                        tuple(w.shape),
                        tuple(target_mod.weight.shape),
                    )

                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))
                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )
                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                w_norm = float(torch.norm(target_mod.weight.data).item())
                if w_norm < 1e-4:
                    logger.warning("OHA: weight '%s' has near-zero norm (%.2e)", proj_key, w_norm)

        # Sanity check (OHA-FIX-13)
        try:
            with torch.no_grad():
                d = next(self.parameters()).device
                dt = self.query.weight.dtype
                dummy = torch.randn(1, 4, self.dim, device=d, dtype=dt)
                _ = self.forward(dummy)
        except Exception as e:
            logger.debug("OHA forward sanity check warning: %s", e)

    # ---- mask helper (OHA-FIX-10) ----
    @staticmethod
    def _prepare_valid_mask(attention_mask, B, N, device):
        if attention_mask is None:
            return torch.ones(B, N, device=device, dtype=torch.bool)
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 3:
            m = attention_mask
            if m.dtype == torch.bool:
                return m.squeeze(1)
            if m.dtype.is_floating_point:
                return (m > -1e4).squeeze(1)
            return (m > 0).squeeze(1)
        if attention_mask.dim() == 4:
            m = attention_mask
            if m.dtype == torch.bool:
                v = m
            elif m.dtype.is_floating_point:
                v = m > -1e4
            else:
                v = m > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ---- forward ----
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        position_embeddings=None,
        past_key_value=None,
        layer_idx=0,
        output_attentions=False,
        *args,
        **kwargs,
    ):
        B, N, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = self.query(hidden_states).view(B, N, H, d).permute(0, 2, 1, 3)
        K = self.key(hidden_states).view(B, N, self.num_kv_heads, d).permute(0, 2, 1, 3)
        V = self.value(hidden_states).view(B, N, self.num_kv_heads, d).permute(0, 2, 1, 3)

        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)

        valid_mask_2d = self._prepare_valid_mask(attention_mask, B, N, device)  # (B, N) bool
        key_valid_4d = valid_mask_2d.unsqueeze(1).unsqueeze(2)  # (B,1,1,N)

        scale = 1.0 / math.sqrt(d)
        idx = torch.arange(N, device=device)
        diff = idx.unsqueeze(1) - idx.unsqueeze(0)  # (N, N)

        # ============================================================
        # 1. Local sliding-window attention
        # ============================================================
        if self.is_causal:
            band_mask = (diff >= 0) & (diff <= self.window_radius)
        else:
            band_mask = diff.abs() <= self.window_radius
        local_mask = band_mask.unsqueeze(0).unsqueeze(0) & key_valid_4d

        local_scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale
        local_scores = local_scores.masked_fill(~local_mask, -1e4)
        local_scores = torch.nan_to_num(local_scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        local_probs = F.softmax(local_scores, dim=-1).to(dtype)
        local_probs = torch.nan_to_num(local_probs, nan=0.0)
        local_probs = self.dropout(local_probs)
        local_ctx = torch.matmul(local_probs, V)  # (B, H, N, d)

        # ============================================================
        # 2. Global compressed attention (OHA-FIX-03, OHA-FIX-14)
        # ============================================================
        R = self.compression_rank
        chunk_size = max(1, (N + R - 1) // R)

        # Slot layout: slot r covers [r·c, min((r+1)·c - 1, N-1)]
        slot_idx = torch.arange(R, device=device)
        chunk_starts = slot_idx * chunk_size  # (R,)
        chunk_ends = torch.minimum(
            (slot_idx + 1) * chunk_size - 1, torch.tensor(N - 1, device=device)
        )  # (R,)
        slot_valid = chunk_starts < N  # (R,) bool

        # comp_k_raw, comp_v_raw shape (B, N, R)
        comp_k_raw = self.compress_k(hidden_states)
        comp_v_raw = self.compress_v(hidden_states)

        # Mask padding positions to -1e4 BEFORE softmax
        pad_mask_3d = (~valid_mask_2d).unsqueeze(-1)  # (B, N, 1)
        comp_k_raw = comp_k_raw.masked_fill(pad_mask_3d, -1e4)
        comp_v_raw = comp_v_raw.masked_fill(pad_mask_3d, -1e4)

        # OHA-FIX-03: chunk-restricted pooling.
        # comp_chunk_mask[n, r] = True iff chunk_starts[r] <= n <= chunk_ends[r]
        pos_idx = torch.arange(N, device=device)  # (N,)
        comp_chunk_mask = (pos_idx.unsqueeze(1) >= chunk_starts.unsqueeze(0)) & (
            pos_idx.unsqueeze(1) <= chunk_ends.unsqueeze(0)
        )  # (N, R)
        comp_chunk_mask = comp_chunk_mask.unsqueeze(0)  # (1, N, R)
        comp_k_raw = comp_k_raw.masked_fill(~comp_chunk_mask, -1e4)
        comp_v_raw = comp_v_raw.masked_fill(~comp_chunk_mask, -1e4)

        comp_k_weights = F.softmax(comp_k_raw.float(), dim=1).to(dtype)  # (B, N, R)
        comp_v_weights = F.softmax(comp_v_raw.float(), dim=1).to(dtype)

        # OHA-FIX-14: zero-out invalid slots (empty chunks when N < R).
        # Note: slot_valid is (R,). Broadcasting to (1, 1, R).
        slot_valid_b = slot_valid.view(1, 1, R)
        comp_k_weights = comp_k_weights * slot_valid_b.to(comp_k_weights.dtype)
        comp_v_weights = comp_v_weights * slot_valid_b.to(comp_v_weights.dtype)

        # ck, cv: (B, H, R, d)
        ck = torch.einsum("bhnd,bnr->bhrd", K, comp_k_weights)
        cv = torch.einsum("bhnd,bnr->bhrd", V, comp_v_weights)

        # Query-global attention
        global_scores = torch.matmul(Q.float(), ck.float().transpose(-2, -1).contiguous()) * scale
        global_scores = torch.nan_to_num(global_scores, nan=-1e4, posinf=1e4, neginf=-1e4)

        if self.is_causal:
            # OHA-FIX-03: query at i sees slot r only if the entire chunk r ends at or before i.
            # chunk_ends[r] <= i
            query_pos = torch.arange(N, device=device).unsqueeze(1)  # (N, 1)
            chunk_ends_b = chunk_ends.unsqueeze(0)  # (1, R)
            global_causal = (query_pos >= chunk_ends_b) & slot_valid.unsqueeze(0)  # (N, R)
            global_scores = global_scores.masked_fill(
                ~global_causal.unsqueeze(0).unsqueeze(0), -1e4
            )

        global_probs = F.softmax(global_scores, dim=-1).to(dtype)
        global_probs = torch.nan_to_num(global_probs, nan=0.0)
        global_probs = self.dropout(global_probs)
        global_ctx = torch.matmul(global_probs, cv)  # (B, H, N, d)

        # ============================================================
        # 3. Gated fusion
        # ============================================================
        if self.use_gate and self.gate is not None:
            gate = torch.sigmoid(self.gate(hidden_states))  # (B, N, H)
            gate = gate.permute(0, 2, 1).unsqueeze(-1)  # (B, H, N, 1)
            context = gate * local_ctx + (1.0 - gate) * global_ctx
        else:
            context = 0.5 * local_ctx + 0.5 * global_ctx

        # 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)

        # Zero-out padded query positions
        context = context * valid_mask_2d.unsqueeze(1).unsqueeze(-1).to(context.dtype)

        context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_proj(context)

        return out, (local_probs if output_attentions else None)

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class PADReAttention (config: Any | None = None,
order: int = 2,
temperature: float = 1.0,
pretrained_attn: nn.Module | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("padre")
@ATTENTION_REGISTRY.register("padre_attention")
class PADReAttention(BaseAttention):
    """
    Polynomial Attention with Dropout Regularization (PADRe).
    Replaces softmax with centered, numerically stable polynomial kernel (1 + S/order)^order.
    """

    def __init__(
        self,
        config: Any | None = None,
        order: int = 2,
        temperature: float = 1.0,
        pretrained_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank=order)
        self.config = config
        self.comp_config = comp_config or _cfg_get(config, "compression_config", None)

        self.hidden_size = _cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768))
        if (
            (config is None or not hasattr(config, "hidden_size"))
            and "hidden_size" not in kwargs
            and pretrained_attn is not None
        ):
            for cand in ["q_lin", "query", "q_proj", "c_attn"]:
                mod = getattr(pretrained_attn, cand, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    self.hidden_size = w.shape[1] if w.ndim == 2 else self.hidden_size
                    break

        self.num_heads = _cfg_get(
            config, "num_attention_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads = _cfg_get(
            config,
            "num_key_value_heads",
            kwargs.get("num_key_value_heads", self.num_heads),
        )
        self.head_dim = _cfg_get(
            config,
            "head_dim",
            kwargs.get(
                "head_dim",
                self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
            ),
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        self.order = _cfg_get(
            self.comp_config,
            "padre_order",
            _cfg_get(config, "padre_order", kwargs.get("order", order)),
        )
        self.temperature = _cfg_get(
            self.comp_config,
            "padre_temperature",
            _cfg_get(config, "padre_temperature", kwargs.get("temperature", temperature)),
        )
        self.max_error = float(
            _cfg_get(self.comp_config, "max_error", kwargs.get("max_error", 0.20))
        )

        # Causal decoder detection
        is_dec = _cfg_get(config, "is_decoder", kwargs.get("is_decoder", False))
        model_type = str(_cfg_get(config, "model_type", "")).lower()
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if any(_mt.startswith(p) for p in _decoder_prefixes):
            is_dec = True
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False
        self.effective_parameter_ratio = 1.0

        # LoRA config
        self.lora_rank = _cfg_get(
            self.comp_config,
            "lora_rank",
            _cfg_get(config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha = _cfg_get(
            self.comp_config,
            "lora_alpha",
            _cfg_get(config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base = self.lora_rank > 0

        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=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.hidden_size,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.hidden_size,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        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=self.freeze_base,
        )

        # Aliases
        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        dropout_p = _cfg_get(config, "attention_probs_dropout_prob", kwargs.get("dropout", 0.1))
        self.dropout = nn.Dropout(dropout_p)

        self.stats: dict[str, Any] = {
            "order": self.order,
            "temperature": self.temperature,
            "lora_rank": self.lora_rank,
            "lora_alpha": self.lora_alpha,
            "num_heads": self.num_heads,
            "num_kv_heads": self.num_kv_heads,
            "is_decoder": self.is_decoder,
            "is_nested_encoder": self.is_nested_encoder,
            "effective_parameter_ratio": self.effective_parameter_ratio,
        }

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if source_attn is not None:
            self.copy_weights(source_attn)

    def copy_weights(self, src: Any) -> None:
        """
        Robust weight copying across GPT-2, RoBERTa/BERT nested, Qwen/LLaMA, and DistilBERT.
        Validates Frobenius norm >= 1e-4.
        """
        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),
        }

        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"],
        }

        # 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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (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])
                    ):
                        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. RoBERTa nested encoder
        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:
                    base_m = getattr(mod, "base_layer", mod)
                    w = getattr(base_m, "weight", getattr(mod, "weight", None))
                    b = getattr(base_m, "bias", getattr(mod, "bias", None))
                    if w is not None:
                        extracted[proj_key] = (
                            w.data.clone(),
                            b.data.clone() if b is not None else None,
                        )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                out_mod = src.output.dense
                base_o = getattr(out_mod, "base_layer", out_mod)
                w = getattr(base_o, "weight", getattr(out_mod, "weight", None))
                b = getattr(base_o, "bias", getattr(out_mod, "bias", None))
                if w is not None:
                    extracted["o"] = (
                        w.data.clone(),
                        b.data.clone() if b is not None else None,
                    )

        # 3. Flat modules (DistilBERT, Qwen, LLaMA)
        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:
                            base_m = getattr(mod, "base_layer", mod)
                            w = getattr(base_m, "weight", getattr(mod, "weight", None))
                            b = getattr(base_m, "bias", getattr(mod, "bias", None))
                            if w is not None:
                                extracted[proj_key] = (
                                    w.data.clone(),
                                    b.data.clone() if b 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:
            logger.warning(
                "PADReAttention copy_weights failed: missing mandatory projections %s from %s.",
                missing_after,
                src.__class__.__name__,
            )
            return None

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        # Check source for LoRA adapters to preserve if present
        source_lora: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {}
        for proj_key, target_mod in target_map.items():
            mod = None
            if proj_key == "o" and hasattr(src, "output") and hasattr(src.output, "dense"):
                mod = src.output.dense
            elif hasattr(src, "self"):
                sub_map = {"q": "query", "k": "key", "v": "value"}
                mod = getattr(src.self, sub_map.get(proj_key, ""), None)
            else:
                for name in proj_candidates.get(proj_key, []):
                    if hasattr(src, name):
                        mod = getattr(src, name)
                        break
            if mod is not None:
                lA = getattr(mod, "lora_A", None)
                lB = getattr(mod, "lora_B", None)
                if lA is not None and lB is not None:
                    # In some implementations lora_A is a Parameter, in PEFT it's a ModuleDict/Linear
                    tA = getattr(lA, "weight", lA)
                    tB = getattr(lB, "weight", lB)
                    if isinstance(tA, (torch.Tensor, nn.Parameter)) and isinstance(
                        tB, (torch.Tensor, nn.Parameter)
                    ):
                        source_lora[proj_key] = (tA.data.clone(), tB.data.clone())

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue

                src_norm = torch.linalg.norm(w.float()).item()
                if src_norm < 1e-4:
                    logger.warning(
                        f"Source weight for {proj_key} has near-zero norm {src_norm:.6f} < 1e-4"
                    )

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                    and w.shape[1] == target_mod.weight.shape[1]
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = torch.linalg.norm(target_mod.weight.float()).item()
                if fnorm < 1e-4:
                    logger.warning(
                        "Compressor copy_weights warning. Proceeding with initialized weights."
                    )

                # Copy source LoRA adapter weights if compatible
                if (
                    proj_key in source_lora
                    and target_mod.lora_A is not None
                    and target_mod.lora_B is not None
                ):
                    slA, slB = source_lora[proj_key]
                    if slA.shape == target_mod.lora_A.shape:
                        target_mod.lora_A.copy_(slA)
                    if slB.shape == target_mod.lora_B.shape:
                        target_mod.lora_B.copy_(slB)

        if self.freeze_base:
            for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
                mod.freeze_base_weights()

        return True

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        Forward pass with polynomial attention:
        1. Linear projection to Q, K, V
        2. Optional RoPE embedding
        3. GQA head expansion
        4. Scaled polynomial scoring with max-centering for stability
        5. Causal and padding masking
        6. Context readout and projection
        """
        B, N, _ = hidden_states.shape

        q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # Apply RoPE if provided
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # GQA head expansion
        if self.num_heads != self.num_kv_heads:
            k = repeat_kv(k, self.num_heads // self.num_kv_heads)
            v = repeat_kv(v, self.num_heads // self.num_kv_heads)

        # Raw scores: (B, H, N, N)
        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1).contiguous()) * scale

        # Determine invalid mask for padding and causal
        invalid_mask = torch.zeros(B, 1, N, N, device=scores.device, dtype=torch.bool)
        if attention_mask is not None:
            valid_tok = get_token_valid_mask(attention_mask, B, N, scores.device)  # (B, N) bool
            invalid_pad = ~valid_tok.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, N)
            invalid_mask = invalid_mask | invalid_pad

        if self.is_decoder:
            causal_mask = torch.tril(torch.ones(N, N, device=scores.device, dtype=torch.bool))
            invalid_causal = ~causal_mask.unsqueeze(0).unsqueeze(0)  # (1, 1, N, N)
            invalid_mask = invalid_mask | invalid_causal

        # Apply temperature and re-center over valid scores for numerical stability
        fill_val = -1e4 if scores.dtype in (torch.float16, torch.bfloat16) else -1e9
        masked_for_max = scores.masked_fill(invalid_mask, fill_val)
        max_scores = masked_for_max.max(dim=-1, keepdim=True).values
        # If all keys are masked, avoid large negative value producing NaN
        thresh = -5e3 if scores.dtype in (torch.float16, torch.bfloat16) else -1e8
        max_scores = torch.where(max_scores <= thresh, torch.zeros_like(max_scores), max_scores)

        scores_centered = (scores - max_scores) / max(self.temperature, 1e-6)

        # Polynomial kernel: (1 + S/order)^order, clamped to avoid negative base or overflow
        base = torch.clamp(1.0 + (scores_centered / self.order), min=0.0, max=10.0)
        attn_weights = base**self.order

        # Zero out invalid positions
        attn_weights = attn_weights.masked_fill(invalid_mask, 0.0)

        # Normalize across key dimension
        denom = attn_weights.sum(dim=-1, keepdim=True) + 1e-8
        attn_weights = attn_weights / denom

        if head_mask is not None:
            hm = head_mask.to(scores.device)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(-1)
            attn_weights = attn_weights * hm

        attn_weights = self.dropout(attn_weights)

        context = torch.matmul(attn_weights, v)
        context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)

        return (output, attn_weights if output_attentions else None)

    def get_ortho_loss(self) -> Any:
        return torch.tensor(0.0, device=self.out_proj.weight.device)

Polynomial Attention with Dropout Regularization (PADRe). Replaces softmax with centered, numerically stable polynomial kernel (1 + S/order)^order.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• BaseAttention
• torch.nn.modules.module.Module

Methods

def copy_weights(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    """
    Robust weight copying across GPT-2, RoBERTa/BERT nested, Qwen/LLaMA, and DistilBERT.
    Validates Frobenius norm >= 1e-4.
    """
    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),
    }

    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"],
    }

    # 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()
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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()
                if (
                    (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])
                ):
                    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. RoBERTa nested encoder
    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:
                base_m = getattr(mod, "base_layer", mod)
                w = getattr(base_m, "weight", getattr(mod, "weight", None))
                b = getattr(base_m, "bias", getattr(mod, "bias", None))
                if w is not None:
                    extracted[proj_key] = (
                        w.data.clone(),
                        b.data.clone() if b is not None else None,
                    )
        if hasattr(src, "output") and hasattr(src.output, "dense"):
            out_mod = src.output.dense
            base_o = getattr(out_mod, "base_layer", out_mod)
            w = getattr(base_o, "weight", getattr(out_mod, "weight", None))
            b = getattr(base_o, "bias", getattr(out_mod, "bias", None))
            if w is not None:
                extracted["o"] = (
                    w.data.clone(),
                    b.data.clone() if b is not None else None,
                )

    # 3. Flat modules (DistilBERT, Qwen, LLaMA)
    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:
                        base_m = getattr(mod, "base_layer", mod)
                        w = getattr(base_m, "weight", getattr(mod, "weight", None))
                        b = getattr(base_m, "bias", getattr(mod, "bias", None))
                        if w is not None:
                            extracted[proj_key] = (
                                w.data.clone(),
                                b.data.clone() if b 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:
        logger.warning(
            "PADReAttention copy_weights failed: missing mandatory projections %s from %s.",
            missing_after,
            src.__class__.__name__,
        )
        return None

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.out_proj,
    }

    # Check source for LoRA adapters to preserve if present
    source_lora: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {}
    for proj_key, target_mod in target_map.items():
        mod = None
        if proj_key == "o" and hasattr(src, "output") and hasattr(src.output, "dense"):
            mod = src.output.dense
        elif hasattr(src, "self"):
            sub_map = {"q": "query", "k": "key", "v": "value"}
            mod = getattr(src.self, sub_map.get(proj_key, ""), None)
        else:
            for name in proj_candidates.get(proj_key, []):
                if hasattr(src, name):
                    mod = getattr(src, name)
                    break
        if mod is not None:
            lA = getattr(mod, "lora_A", None)
            lB = getattr(mod, "lora_B", None)
            if lA is not None and lB is not None:
                # In some implementations lora_A is a Parameter, in PEFT it's a ModuleDict/Linear
                tA = getattr(lA, "weight", lA)
                tB = getattr(lB, "weight", lB)
                if isinstance(tA, (torch.Tensor, nn.Parameter)) and isinstance(
                    tB, (torch.Tensor, nn.Parameter)
                ):
                    source_lora[proj_key] = (tA.data.clone(), tB.data.clone())

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                    continue

            src_norm = torch.linalg.norm(w.float()).item()
            if src_norm < 1e-4:
                logger.warning(
                    f"Source weight for {proj_key} has near-zero norm {src_norm:.6f} < 1e-4"
                )

            if w.shape == target_mod.weight.shape:
                target_mod.weight.copy_(w)
            elif w.t().contiguous().shape == target_mod.weight.shape:
                target_mod.weight.copy_(w.t().contiguous())
            elif (
                w.shape[0] < target_mod.weight.shape[0]
                and target_mod.weight.shape[0] % w.shape[0] == 0
                and w.shape[1] == target_mod.weight.shape[1]
            ):
                repeats = target_mod.weight.shape[0] // w.shape[0]
                target_mod.weight.copy_(w.repeat(repeats, 1))
            else:
                min_0 = min(target_mod.weight.shape[0], w.shape[0])
                min_1 = min(target_mod.weight.shape[1], w.shape[1])
                target_mod.weight.zero_()
                target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

            if b is not None and target_mod.bias is not None:
                if b.shape == target_mod.bias.shape:
                    target_mod.bias.copy_(b)
                elif (
                    b.shape[0] < target_mod.bias.shape[0]
                    and target_mod.bias.shape[0] % b.shape[0] == 0
                ):
                    repeats = target_mod.bias.shape[0] // b.shape[0]
                    target_mod.bias.copy_(b.repeat(repeats))
                else:
                    min_b = min(target_mod.bias.shape[0], b.shape[0])
                    target_mod.bias.zero_()
                    target_mod.bias[:min_b].copy_(b[:min_b])
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

            fnorm = torch.linalg.norm(target_mod.weight.float()).item()
            if fnorm < 1e-4:
                logger.warning(
                    "Compressor copy_weights warning. Proceeding with initialized weights."
                )

            # Copy source LoRA adapter weights if compatible
            if (
                proj_key in source_lora
                and target_mod.lora_A is not None
                and target_mod.lora_B is not None
            ):
                slA, slB = source_lora[proj_key]
                if slA.shape == target_mod.lora_A.shape:
                    target_mod.lora_A.copy_(slA)
                if slB.shape == target_mod.lora_B.shape:
                    target_mod.lora_B.copy_(slB)

    if self.freeze_base:
        for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
            mod.freeze_base_weights()

    return True

Robust weight copying across GPT-2, RoBERTa/BERT nested, Qwen/LLaMA, and DistilBERT. Validates Frobenius norm >= 1e-4.

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """
    Forward pass with polynomial attention:
    1. Linear projection to Q, K, V
    2. Optional RoPE embedding
    3. GQA head expansion
    4. Scaled polynomial scoring with max-centering for stability
    5. Causal and padding masking
    6. Context readout and projection
    """
    B, N, _ = hidden_states.shape

    q = (
        self.q_proj(hidden_states)
        .view(B, N, self.num_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    k = (
        self.k_proj(hidden_states)
        .view(B, N, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v = (
        self.v_proj(hidden_states)
        .view(B, N, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # Apply RoPE if provided
    if position_embeddings is None:
        position_embeddings = kwargs.get("position_embeddings")
    if position_embeddings is not None:
        cos, sin = position_embeddings
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    # GQA head expansion
    if self.num_heads != self.num_kv_heads:
        k = repeat_kv(k, self.num_heads // self.num_kv_heads)
        v = repeat_kv(v, self.num_heads // self.num_kv_heads)

    # Raw scores: (B, H, N, N)
    scale = 1.0 / math.sqrt(self.head_dim)
    scores = torch.matmul(q, k.transpose(-2, -1).contiguous()) * scale

    # Determine invalid mask for padding and causal
    invalid_mask = torch.zeros(B, 1, N, N, device=scores.device, dtype=torch.bool)
    if attention_mask is not None:
        valid_tok = get_token_valid_mask(attention_mask, B, N, scores.device)  # (B, N) bool
        invalid_pad = ~valid_tok.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, N)
        invalid_mask = invalid_mask | invalid_pad

    if self.is_decoder:
        causal_mask = torch.tril(torch.ones(N, N, device=scores.device, dtype=torch.bool))
        invalid_causal = ~causal_mask.unsqueeze(0).unsqueeze(0)  # (1, 1, N, N)
        invalid_mask = invalid_mask | invalid_causal

    # Apply temperature and re-center over valid scores for numerical stability
    fill_val = -1e4 if scores.dtype in (torch.float16, torch.bfloat16) else -1e9
    masked_for_max = scores.masked_fill(invalid_mask, fill_val)
    max_scores = masked_for_max.max(dim=-1, keepdim=True).values
    # If all keys are masked, avoid large negative value producing NaN
    thresh = -5e3 if scores.dtype in (torch.float16, torch.bfloat16) else -1e8
    max_scores = torch.where(max_scores <= thresh, torch.zeros_like(max_scores), max_scores)

    scores_centered = (scores - max_scores) / max(self.temperature, 1e-6)

    # Polynomial kernel: (1 + S/order)^order, clamped to avoid negative base or overflow
    base = torch.clamp(1.0 + (scores_centered / self.order), min=0.0, max=10.0)
    attn_weights = base**self.order

    # Zero out invalid positions
    attn_weights = attn_weights.masked_fill(invalid_mask, 0.0)

    # Normalize across key dimension
    denom = attn_weights.sum(dim=-1, keepdim=True) + 1e-8
    attn_weights = attn_weights / denom

    if head_mask is not None:
        hm = head_mask.to(scores.device)
        if hm.dim() == 1:
            hm = hm.view(1, -1, 1, 1)
        elif hm.dim() == 2:
            hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
        elif hm.dim() == 3:
            hm = hm.unsqueeze(-1)
        attn_weights = attn_weights * hm

    attn_weights = self.dropout(attn_weights)

    context = torch.matmul(attn_weights, v)
    context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)

    if self.is_nested_encoder:
        output = context
    else:
        output = self.out_proj(context)

    return (output, attn_weights if output_attentions else None)

Forward pass with polynomial attention: 1. Linear projection to Q, K, V 2. Optional RoPE embedding 3. GQA head expansion 4. Scaled polynomial scoring with max-centering for stability 5. Causal and padding masking 6. Context readout and projection

def get_ortho_loss(self) ‑> Any

Expand source code

def get_ortho_loss(self) -> Any:
    return torch.tensor(0.0, device=self.out_proj.weight.device)

Inherited members

• BaseAttention:
  • extra_repr

class PHANTOMAttention (config: Any | None = None,
rank: int | None = 16,
phantom_cfg: PHANTOMConfig | None = None,
pretrained_attn: Any | None = None,
bert_mode: bool = False,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("phantom")
@ATTENTION_REGISTRY.register("phantom_module")
class PHANTOMAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 16,
        phantom_cfg: PHANTOMConfig | None = None,
        pretrained_attn: Any | None = None,
        bert_mode: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank or 16)
        self.config = config
        comp_config = getattr(config, "compression_config", None)

        cur_rank = rank or getattr(
            comp_config,
            "cur_rank",
            getattr(config, "cur_rank", getattr(config, "svd_rank", 16)),
        )
        if phantom_cfg is None:
            phantom_cfg = PHANTOMConfig(cur_rank=cur_rank)
        elif rank is not None:
            phantom_cfg.cur_rank = rank
        self.cfg: PHANTOMConfig = phantom_cfg

        model_config = config or kwargs.get("model_config")
        self.hidden_size: int = getattr(model_config, "hidden_size", kwargs.get("hidden_size", 768))
        self.dim: int = self.hidden_size
        self.num_heads: int = getattr(
            model_config, "num_attention_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads: int = getattr(
            model_config,
            "num_key_value_heads",
            kwargs.get("num_key_value_heads", self.num_heads),
        )
        self.head_dim: int = getattr(
            model_config,
            "head_dim",
            kwargs.get("head_dim", self.dim // self.num_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
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa: bool = True

        is_dec = getattr(config, "is_decoder", kwargs.get("is_decoder", False))
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        target_dtype = kwargs.get("dtype")
        if target_dtype is None and source_attn is not None:
            target_dtype = next(source_attn.parameters(), torch.tensor([])).dtype
            if not isinstance(target_dtype, torch.dtype):
                target_dtype = torch.float32

        # Projections
        self.q_proj = FactorisedLinear_CUR(
            self.dim,
            self.all_head_size,
            cur_rank=self.cfg.cur_rank,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.k_proj = FactorisedLinear_CUR(
            self.dim,
            self.all_kv_head_size,
            cur_rank=self.cfg.cur_rank,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.v_proj = FactorisedLinear_CUR(
            self.dim,
            self.all_kv_head_size,
            cur_rank=self.cfg.cur_rank,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.out_proj = FactorisedLinear_CUR(
            self.all_head_size,
            self.dim,
            cur_rank=self.cfg.cur_rank,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )

        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        self.dropout = nn.Dropout(self.cfg.dropout)

        if source_attn is not None:
            self.copy_weights(source_attn)

        if target_dtype is not None:
            self.to(target_dtype)

    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")
            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
        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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (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])
                    ):
                        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. 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
        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:
            logger.warning(
                f"PHANTOMAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}."
            )
            return

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        target_mod.copy_cur(torch.eye(self.all_head_size, self.dim), None)
                        continue
                    else:
                        logger.warning(f"Missing required weights for {proj_key}")
                        return

                if w.shape != (target_mod.out_features, target_mod.in_features):
                    if w.t().contiguous().shape == (
                        target_mod.out_features,
                        target_mod.in_features,
                    ):
                        w = w.t().contiguous()
                    elif (
                        w.shape[0] < target_mod.out_features
                        and target_mod.out_features % w.shape[0] == 0
                    ):
                        repeats = target_mod.out_features // w.shape[0]
                        w = w.repeat(repeats, 1)
                    else:
                        full_w = torch.zeros(
                            (target_mod.out_features, target_mod.in_features),
                            dtype=w.dtype,
                            device=w.device,
                        )
                        min_0 = min(target_mod.out_features, w.shape[0])
                        min_1 = min(target_mod.in_features, w.shape[1])
                        full_w[:min_0, :min_1].copy_(w[:min_0, :min_1])
                        w = full_w

                target_mod.copy_cur(w, b)

                fnorm = torch.linalg.norm(target_mod.weight.float()).item()
                if fnorm < 1e-4:
                    logger.warning(f"Weight copy notice for {proj_key}: norm {fnorm:.6f} < 1e-4")

        if self.freeze_base:
            for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
                mod.freeze_base_weights()

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim

        Q = self.q_proj(hidden_states).view(B, N, H, D).permute(0, 2, 1, 3)
        K = self.k_proj(hidden_states).view(B, N, self.num_kv_heads, D).permute(0, 2, 1, 3)
        V = self.v_proj(hidden_states).view(B, N, self.num_kv_heads, D).permute(0, 2, 1, 3)

        # Apply RoPE if provided
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # Repeat KV for GQA
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            K = repeat_kv(K, repeats)
            V = repeat_kv(V, repeats)

        token_valid = get_token_valid_mask(attention_mask, B, N, hidden_states.device)
        is_causal = self.is_decoder or kwargs.get("is_causal", False)

        context = lsh_attention(
            Q,
            K,
            V,
            rounds=self.cfg.lsh_rounds,
            bucket_size=self.cfg.lsh_bucket_size,
            softpick_temp=(self.cfg.softpick_temperature if self.cfg.softpick_enabled else 1.0),
            dropout=self.dropout if self.cfg.softpick_enabled else None,
            token_valid=token_valid,
            is_causal=is_causal,
        )

        # head_mask handling
        if head_mask is not None:
            if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                context = context * head_mask.view(1, self.num_heads, 1, 1).to(context.dtype)
            elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                context = context * head_mask.view(head_mask.size(0), self.num_heads, 1, 1).to(
                    context.dtype
                )
            elif head_mask.dim() == 4:
                context = context * head_mask.to(context.dtype)

        context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            return context, None

        output = self.out_proj(context)
        return output, None

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class PHANTOMConfig (cur_rank: int = 16,
lsh_rounds: int = 4,
lsh_bucket_size: int = 4,
softpick_enabled: bool = True,
softpick_temperature: float = 1.0,
dropout: float = 0.1)

Expand source code

@dataclass
class PHANTOMConfig:
    cur_rank: int = 16
    lsh_rounds: int = 4
    lsh_bucket_size: int = 4
    softpick_enabled: bool = True
    softpick_temperature: float = 1.0
    dropout: float = 0.1

PHANTOMConfig(cur_rank: 'int' = 16, lsh_rounds: 'int' = 4, lsh_bucket_size: 'int' = 4, softpick_enabled: 'bool' = True, softpick_temperature: 'float' = 1.0, dropout: 'float' = 0.1)

Instance variables

var cur_rank : int

The type of the None singleton.

var dropout : float

The type of the None singleton.

var lsh_bucket_size : int

The type of the None singleton.

var lsh_rounds : int

The type of the None singleton.

var softpick_enabled : bool

The type of the None singleton.

var softpick_temperature : float

The type of the None singleton.

class PerformerAttention (config: Any | None = None,
rank: int | None = None,
eps: float = 1e-06,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
is_nested_encoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "performer",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "performer_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class PerformerAttention(BaseAttention):
    is_custom_attention: bool = True
    triton_kernel = "linear_attn"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        eps: float = 1e-6,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        comp_config: Any | None = None,
        is_nested_encoder: bool | None = None,  # FIX 4: explicit override
        **kwargs: Any,
    ) -> None:
        source_attn = pretrained_attn if pretrained_attn is not None else old_attn
        super().__init__(config, rank)

        # ---------------- Hidden size & heads ----------------
        hidden_size = kwargs.get("hidden_size", kwargs.get("dim", kwargs.get("d_model")))
        num_heads = kwargs.get(
            "num_attention_heads", kwargs.get("num_heads", kwargs.get("n_heads"))
        )
        num_kv_heads = kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads"))

        if hidden_size is None or num_heads is None:
            if config is not None:
                if hidden_size is None:
                    hidden_size = _cfg_get(
                        config,
                        "hidden_size",
                        _cfg_get(config, "dim", _cfg_get(config, "d_model", 768)),
                    )
                if num_heads is None:
                    num_heads = _cfg_get(
                        config,
                        "num_attention_heads",
                        _cfg_get(config, "n_heads", _cfg_get(config, "num_heads", 12)),
                    )
                if num_kv_heads is None:
                    num_kv_heads = _cfg_get(
                        config,
                        "num_key_value_heads",
                        _cfg_get(config, "num_kv_heads", num_heads),
                    )
            elif source_attn is not None:
                if hidden_size is None:
                    out_layer = getattr(
                        source_attn,
                        "out_lin",
                        getattr(
                            source_attn,
                            "dense",
                            getattr(
                                source_attn,
                                "c_proj",
                                getattr(source_attn, "o_proj", None),
                            ),
                        ),
                    )
                    hidden_size = getattr(
                        out_layer,
                        "out_features",
                        getattr(
                            source_attn,
                            "hidden_size",
                            getattr(source_attn, "embed_dim", 768),
                        ),
                    )
                if num_heads is None:
                    num_heads = getattr(
                        source_attn,
                        "num_attention_heads",
                        getattr(
                            source_attn,
                            "n_heads",
                            getattr(source_attn, "num_heads", 12),
                        ),
                    )
                if num_kv_heads is None:
                    num_kv_heads = getattr(
                        source_attn,
                        "num_key_value_heads",
                        getattr(source_attn, "num_kv_heads", num_heads),
                    )
            else:
                hidden_size = hidden_size if hidden_size is not None else 768
                num_heads = num_heads if num_heads is not None else 12
                num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads

        if num_kv_heads is None:
            num_kv_heads = num_heads

        # FIX 8: validate divisibility
        if num_heads <= 0 or num_kv_heads <= 0:
            raise ValueError(
                f"Invalid head counts: num_heads={num_heads}, num_kv_heads={num_kv_heads}"
            )
        if num_heads % num_kv_heads != 0:
            raise ValueError(
                f"num_heads ({num_heads}) must be divisible by num_kv_heads ({num_kv_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 = self.dim // self.num_heads if self.num_heads > 0 else 64
        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.eps = eps

        # ---------------- num_features (FIX 7) ----------------
        num_features_default = rank if rank is not None else 64
        self.num_features = _resolve_config_value(
            config,
            comp_config,
            kwargs,
            names=("performer_rank", "rank", "num_features"),
            default=num_features_default,
        )
        # FIX 14: ensure at least 1
        self.num_features = max(1, int(self.num_features))

        # ---------------- kernel_type ----------------
        self.kernel_type = _resolve_config_value(
            config,
            comp_config,
            kwargs,
            names=("performer_kernel_type", "kernel_type"),
            default="favor_plus",
        )

        # ---------------- Decoder detection ----------------
        model_type = getattr(config, "model_type", "").lower() if config else ""
        is_dec = (
            getattr(
                config,
                "is_decoder",
                getattr(config, "is_causal", kwargs.get("is_causal", False)),
            )
            if config
            else False
        )
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(_mt.startswith(p) for p in _decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)

        # ---------------- LoRA rank & alpha (FIX 7) ----------------
        self.lora_rank = int(
            _resolve_config_value(
                config,
                comp_config,
                kwargs,
                names=("lora_rank",),
                default=16,
            )
        )
        self.lora_alpha = float(
            _resolve_config_value(
                config,
                comp_config,
                kwargs,
                names=("lora_alpha",),
                default=32.0,
            )
        )
        self.freeze_base = self.lora_rank > 0

        # ---------------- Nested encoder detection (FIX 4) ----------------
        # If explicit value passed, use it. Otherwise infer from source_attn.
        if is_nested_encoder is not None:
            self.is_nested_encoder = bool(is_nested_encoder)
        else:
            self.is_nested_encoder = False  # will be updated by copy_weights

        # ---------------- Projections (FIX 2) ----------------
        # out_proj gets LoRA only if NOT nested encoder (since nested encoder
        # skips out_proj in forward pass).
        # We determine nested-ness AFTER copy_weights, but we need to build
        # modules before. So: build with lora on out_proj tentatively; if
        # copy_weights reveals nested, we zero-out / disable it.
        self.query = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.key = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.value = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        # out_proj LoRA is created conditionally. We start with lora_rank,
        # and if copy_weights detects nested, we disable it.
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Aliases (no longer registered as parameters, only attribute access)
        self.q_proj = self.query
        self.k_proj = self.key
        self.v_proj = self.value

        # ---------------- Omega matrix ----------------
        omega = self._build_orthogonal_matrix(self.head_dim, self.num_features)
        self.register_buffer("omega", omega, persistent=True)

        # ---------------- Dropout (FIX 10) ----------------
        if config is not None:
            dropout_p = getattr(
                config,
                "attention_probs_dropout_prob",
                getattr(config, "dropout", 0.1),
            )
        else:
            dropout_p = 0.1
        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(source_attn, "rotary_emb", None)
        self.triton_kwargs = {
            "phi": "identity",
            "eps": self.eps,
            "cos_weight": None,
            "sin_weight": None,
        }

        # ---------------- Copy weights from source ----------------
        if source_attn is not None:
            success = self.copy_weights(source_attn)
            if not success:
                logger.warning(
                    "PerformerAttention: copy_weights failed. Proceeding with initialized weights."
                )

        # FIX 2: after is_nested_encoder is known, disable LoRA on out_proj
        # if nested encoder (since out_proj is not used in forward).
        if self.is_nested_encoder and self.lora_rank > 0:
            # Move out_proj LoRA params to a "dummy" state so they don't
            # participate in optimization. We freeze them and zero-init.
            if self.out_proj.lora_A is not None:
                with torch.no_grad():
                    self.out_proj.lora_A.zero_()
                self.out_proj.lora_A.requires_grad = False
            if self.out_proj.lora_B is not None:
                with torch.no_grad():
                    self.out_proj.lora_B.zero_()
                self.out_proj.lora_B.requires_grad = False
            logger.info(
                "PerformerAttention: is_nested_encoder=True, disabling LoRA "
                "on out_proj (out_proj is bypassed in forward pass)."
            )

    # ------------------------------------------------------------
    def _build_orthogonal_matrix(self, d: int, m: int) -> torch.Tensor:
        blocks = []
        remaining = m
        while remaining > 0:
            block_size = min(d, remaining)
            g = torch.randn(d, block_size)
            q, _ = torch.linalg.qr(g)
            blocks.append(q)
            remaining -= block_size
        W = torch.cat(blocks, dim=1)
        W = W * math.sqrt(d)
        return W.t().contiguous()

    # ------------------------------------------------------------
    def _feature_map(self, x: torch.Tensor) -> torch.Tensor:
        """FAVOR+ feature map.

        FIX 1: removed per-token max_proj subtraction (bias).
        FIX 6: removed redundant +eps.
        FIX 11: removed .contiguous() (matmul handles non-contiguous).
        FIX 12: removed dead clamp(max=10.0).
        """
        # x: [B, H, N, d]
        # omega: [M, d]
        omega_t = self.omega.to(dtype=x.dtype, device=x.device).t()  # [d, M]
        proj = torch.matmul(x, omega_t)  # [B, H, N, M]

        if self.kernel_type in ("favor_plus", "exp", "favor"):
            # FIX 1: no per-token max subtraction
            diag_sq = torch.sum(x**2, dim=-1, keepdim=True) * 0.5  # [B,H,N,1]
            # FIX 12: only clamp min
            diff = torch.clamp(proj - diag_sq, min=-30.0)
            phi = torch.exp(diff)  # FIX 6: no +eps
            phi = phi / math.sqrt(self.num_features)
        elif self.kernel_type == "elu":
            phi = (F.elu(proj) + 1.0) / math.sqrt(self.num_features)
        else:
            phi = (F.relu(proj) + self.eps) / math.sqrt(self.num_features)
        return phi

    # ------------------------------------------------------------
    def copy_weights(self, src: Any) -> bool:
        """Copy weights from source attention module.

        FIX 13: returns bool (was None).
        FIX 4: sets is_nested_encoder robustly.
        """
        if src is None:
            return False

        cls_name = src.__class__.__name__
        has_out_lin = hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj")

        # FIX 4: robust nested detection
        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

        # Only override if we didn't get explicit is_nested_encoder in __init__
        # We know we didn't if current value is False (default) and no source
        # was given... but actually we should always trust the source.
        # Let's just set it (unless caller explicitly passed one, in which
        # case we don't override — but we can't detect that here, so we
        # assume the source is authoritative).
        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 Conv1D ----
        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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (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])
                    ):
                        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. RoBERTa/BERT 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 ----
        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

        # ---- 4. Fallback via UniversalAdapter ----
        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:
                pass

        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:
            logger.warning(
                "PerformerAttention copy_weights failed: missing mandatory projections %s from %s.",
                missing_after,
                src.__class__.__name__,
            )
            return False

        # ---- 5. Copy into target modules ----
        target_map = {
            "q": self.query,
            "k": self.key,
            "v": self.value,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()

                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))

                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device,
                        dtype=target_mod.weight.dtype,
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(
                            device=target_mod.bias.device,
                            dtype=target_mod.bias.dtype,
                        )
                    )
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                w_norm = float(torch.norm(target_mod.weight.data).item())
                if w_norm < 1e-4:
                    logger.warning(
                        "PerformerAttention weight verification notice for %s: "
                        "norm too small (%.2e)",
                        proj_key,
                        w_norm,
                    )

        # ---- 6. Ensure base weights frozen, LoRA A/B trainable ----
        if self.lora_rank > 0:
            self.query.freeze_base_weights()
            self.key.freeze_base_weights()
            self.value.freeze_base_weights()
            if not self.is_nested_encoder:
                self.out_proj.freeze_base_weights()
            # If nested encoder, out_proj LoRA was already disabled

        return True

    # ------------------------------------------------------------
    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,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        output_attentions: bool | None = False,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        # FIX 5: explicit error instead of silent None
        if output_attentions:
            raise NotImplementedError(
                "PerformerAttention does not materialize attention matrices. "
                "FAVOR+ computes attention implicitly and does not store the "
                "O(N^2) attention scores."
            )

        if encoder_hidden_states is not None:
            # Cross-attention not supported
            return hidden_states, None

        while isinstance(hidden_states, (tuple, list)):
            hidden_states = hidden_states[0]

        B, N, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        scale = d**-0.25
        Q = self.query(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous() * scale
        K = (
            self.key(hidden_states).view(B, N, self.num_kv_heads, d).transpose(1, 2).contiguous()
            * scale
        )
        V = self.value(hidden_states).view(B, N, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # ---- RoPE ----
        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 self.rotary_emb is not None:
            pos_ids = kwargs.get("position_ids")
            if pos_ids is None:
                pos_ids = torch.arange(N, device=device).unsqueeze(0).expand(B, -1)
            try:
                cos, sin = self.rotary_emb(V, pos_ids)
                Q, K = apply_rotary_pos_emb(Q, K, cos, sin, pos_ids)
            except Exception as e:
                logger.debug("RoPE application failed: %s", e)

        # ---- GQA expansion ----
        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        # ---- Valid token mask ----
        valid_mask = get_token_valid_mask(attention_mask, B, N, device)
        if valid_mask is not None:
            mask_expanded = valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
        else:
            mask_expanded = torch.ones(B, 1, N, 1, device=device, dtype=dtype)

        # ---- Feature maps ----
        # FIX 3: mask Q_feat in BOTH branches
        Q_feat = self._feature_map(Q) * mask_expanded
        K_feat = self._feature_map(K) * mask_expanded

        # ---- Attention via shared Triton dispatcher ----
        context = triton_dispatch(
            self.triton_kernel,
            Q_feat,
            K_feat,
            V,
            is_causal=self.is_decoder,
            **self.triton_kwargs,
        )

        context = torch.nan_to_num(context, nan=0.0)
        context = self.dropout(context)

        # ---- Head mask (FIX 9) ----
        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.view(head_mask.size(0), -1, 1, 1)
            elif head_mask.dim() == 3:
                hm = head_mask.unsqueeze(-1)
            else:
                hm = head_mask
            context = context * hm

        context = context * mask_expanded
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        # FIX 2: for nested encoder, skip out_proj (LoRA disabled)
        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_proj(context)

        # ---- Shape sanity ----
        if out.shape[-1] != hidden_states.shape[-1]:
            if out.shape[-1] < hidden_states.shape[-1]:
                out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
            else:
                out = out[..., : hidden_states.shape[-1]]

        return out.to(dtype), None

    # ------------------------------------------------------------
    def _forward_causal(
        self,
        Q_feat: torch.Tensor,
        K_feat: torch.Tensor,
        V: torch.Tensor,
        B: int,
        H: int,
        d: int,
        N: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Exact causal FAVOR+ prefix scan."""
        M = Q_feat.shape[-1]
        chunk_size = max(16, min(64, N // 4))
        num_chunks = (N + chunk_size - 1) // chunk_size

        S = torch.zeros(B, H, M, d, dtype=torch.float32, device=device)
        z = torch.zeros(B, H, M, dtype=torch.float32, device=device)
        num_list = []
        den_list = []

        for c in range(num_chunks):
            start = c * chunk_size
            end = min(start + chunk_size, N)
            L = end - start

            q_c = Q_feat[:, :, start:end].float()
            k_c = K_feat[:, :, start:end].float()
            v_c = V[:, :, start:end].float()

            causal_mask = torch.tril(torch.ones(L, L, device=device, dtype=torch.float32))
            intra_attn = torch.matmul(q_c, k_c.transpose(-1, -2)) * causal_mask
            intra_num = torch.matmul(intra_attn, v_c)
            inter_num = torch.matmul(q_c, S)
            num_list.append(intra_num + inter_num)

            intra_den = intra_attn.sum(dim=-1, keepdim=True)
            inter_den = torch.matmul(q_c, z.unsqueeze(-1))
            den_list.append(intra_den + inter_den)

            S = S + torch.matmul(k_c.transpose(-1, -2), v_c)
            z = z + k_c.sum(dim=2)

        num = torch.cat(num_list, dim=2)
        den = torch.cat(den_list, dim=2).clamp(min=self.eps)
        return (num / den).to(dtype)

    # ------------------------------------------------------------
    def _forward_bidirectional(
        self,
        Q_feat: torch.Tensor,
        K_feat: torch.Tensor,
        V: torch.Tensor,
        B: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Standard linear attention (non-causal)."""
        # FIX 9: only upcast to fp32 when needed (fp16 case)
        compute_dtype = torch.float32 if dtype == torch.float16 else dtype
        Kf = K_feat.to(compute_dtype)
        Vf = V.to(compute_dtype)
        Qf = Q_feat.to(compute_dtype)

        KV = torch.matmul(Kf.transpose(-2, -1), Vf)
        K_sum = Kf.sum(dim=-2, keepdim=True)
        num = torch.matmul(Qf, KV)
        den = torch.matmul(Qf, K_sum.transpose(-2, -1)) + self.eps
        return (num / den).to(dtype)

Every 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 triton_kernel

The type of the None singleton.

Methods

def copy_weights(self,
src: Any) ‑> bool

Expand source code

def copy_weights(self, src: Any) -> bool:
    """Copy weights from source attention module.

    FIX 13: returns bool (was None).
    FIX 4: sets is_nested_encoder robustly.
    """
    if src is None:
        return False

    cls_name = src.__class__.__name__
    has_out_lin = hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj")

    # FIX 4: robust nested detection
    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

    # Only override if we didn't get explicit is_nested_encoder in __init__
    # We know we didn't if current value is False (default) and no source
    # was given... but actually we should always trust the source.
    # Let's just set it (unless caller explicitly passed one, in which
    # case we don't override — but we can't detect that here, so we
    # assume the source is authoritative).
    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 Conv1D ----
    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()
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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()
                if (
                    (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])
                ):
                    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. RoBERTa/BERT 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 ----
    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

    # ---- 4. Fallback via UniversalAdapter ----
    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:
            pass

    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:
        logger.warning(
            "PerformerAttention copy_weights failed: missing mandatory projections %s from %s.",
            missing_after,
            src.__class__.__name__,
        )
        return False

    # ---- 5. Copy into target modules ----
    target_map = {
        "q": self.query,
        "k": self.key,
        "v": self.value,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                continue

            if w.ndim != 2:
                w = w.reshape(w.size(0), -1)

            if (
                w.shape != target_mod.weight.shape
                and target_mod.weight.shape == w.t().contiguous().shape
            ):
                w = w.t().contiguous()

            out_d = min(target_mod.weight.size(0), w.size(0))
            in_d = min(target_mod.weight.size(1), w.size(1))

            target_mod.weight.data[:out_d, :in_d].copy_(
                w[:out_d, :in_d].to(
                    device=target_mod.weight.device,
                    dtype=target_mod.weight.dtype,
                )
            )

            if b is not None and target_mod.bias is not None:
                min_b = min(target_mod.bias.numel(), b.numel())
                target_mod.bias.data[:min_b].copy_(
                    b[:min_b].to(
                        device=target_mod.bias.device,
                        dtype=target_mod.bias.dtype,
                    )
                )
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

            w_norm = float(torch.norm(target_mod.weight.data).item())
            if w_norm < 1e-4:
                logger.warning(
                    "PerformerAttention weight verification notice for %s: "
                    "norm too small (%.2e)",
                    proj_key,
                    w_norm,
                )

    # ---- 6. Ensure base weights frozen, LoRA A/B trainable ----
    if self.lora_rank > 0:
        self.query.freeze_base_weights()
        self.key.freeze_base_weights()
        self.value.freeze_base_weights()
        if not self.is_nested_encoder:
            self.out_proj.freeze_base_weights()
        # If nested encoder, out_proj LoRA was already disabled

    return True

Copy weights from source attention module.

FIX 13: returns bool (was None). FIX 4: sets is_nested_encoder robustly.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class PrimalAttention (config: Any | None = None,
rank: int = 64,
pinv_iters: int = 2,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("primal")
@ATTENTION_REGISTRY.register("primal_attention")
class PrimalAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 64,
        pinv_iters: int = 2,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        if rank is None:
            rank = 64
        super().__init__(config, rank)
        self.config = config
        comp_config = getattr(config, "compression_config", None)

        self.dim = getattr(config, "hidden_size", kwargs.get("hidden_size", 768))
        self.num_heads = getattr(
            config, "num_attention_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads = getattr(
            config,
            "num_key_value_heads",
            kwargs.get("num_key_value_heads", self.num_heads),
        )
        self.head_dim = getattr(
            config,
            "head_dim",
            kwargs.get("head_dim", self.dim // self.num_heads),
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        # PR-FIX-03: guard self.rank against None
        _rank_candidate = None
        if comp_config is not None:
            _rank_candidate = getattr(comp_config, "rank", None)
        if _rank_candidate is None:
            _rank_candidate = kwargs.pop("rank", None)
        if _rank_candidate is None:
            _rank_candidate = rank
        self.rank = int(_rank_candidate) if _rank_candidate else 64

        self.pinv_iters = pinv_iters

        is_dec = getattr(config, "is_decoder", kwargs.get("is_decoder", False))
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa = True

        # PR-FIX-07: chunked head processing threshold.
        # If rank is large, split heads into chunks of CHUNK_HEADS to bound memory.
        self.chunk_heads = int(kwargs.get("chunk_heads", 0))  # 0 = disabled

        # Projections
        self.query = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.key = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.value = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.q_proj = self.query
        self.k_proj = self.key
        self.v_proj = self.value

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if source_attn is not None:
            self.copy_weights(source_attn)

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    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
        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[0] < c_w.shape[1])
                )
                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. RoBERTa / BERT 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, GPT-2 style without c_attn)
        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"PrimalAttention copy_weights failed: missing mandatory projections "
                f"{missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.query,
            "k": self.key,
            "v": self.value,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]

                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    else:
                        logger.warning(
                            "PrimalAttention.copy_weights: no weight for %s, "
                            "keeping initialized weights.",
                            proj_key,
                        )
                        continue

                # Shape alignment
                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = torch.linalg.norm(target_mod.weight.float()).item()
                if fnorm < 1e-4:
                    logger.warning(
                        "PrimalAttention.copy_weights: %s has near-zero Frobenius norm (%.2e).",
                        proj_key,
                        fnorm,
                    )

        if self.freeze_base:
            for mod in [self.query, self.key, self.value, self.out_proj]:
                mod.freeze_base_weights()

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, seq_len, _ = hidden_states.shape

        Q = (
            self.query(hidden_states)
            .view(batch_size, seq_len, self.num_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        K = (
            self.key(hidden_states)
            .view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        V = (
            self.value(hidden_states)
            .view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )

        # RoPE
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # GQA expansion
        if self.num_heads != self.num_kv_heads:
            num_reps = self.num_heads // self.num_kv_heads
            K = repeat_kv(K, num_reps)
            V = repeat_kv(V, num_reps)

        # Landmark computation
        num_landmarks = min(self.rank, seq_len)
        if seq_len % num_landmarks == 0:
            seg = seq_len // num_landmarks
            Q_land = Q.reshape(batch_size, self.num_heads, num_landmarks, seg, self.head_dim).mean(
                dim=3
            )
            K_land = K.reshape(batch_size, self.num_heads, num_landmarks, seg, self.head_dim).mean(
                dim=3
            )
        else:
            idx = torch.linspace(
                0,
                seq_len - 1,
                num_landmarks,
                dtype=torch.long,
                device=hidden_states.device,
            )
            Q_land = Q[:, :, idx, :]
            K_land = K[:, :, idx, :]

        scale = 1.0 / (self.head_dim**0.5)

        # PR-FIX-01 + PR-FIX-07: chunked head processing to bound peak memory.
        # The chunk_heads attribute (default 0) disables chunking.
        if self.chunk_heads > 0 and self.num_heads > self.chunk_heads:
            head_chunks = [
                (h0, min(h0 + self.chunk_heads, self.num_heads))
                for h0 in range(0, self.num_heads, self.chunk_heads)
            ]
        else:
            head_chunks = [(0, self.num_heads)]

        context_heads = []
        attn_for_return: torch.Tensor | None = None

        for h0, h1 in head_chunks:
            Qc = Q[:, h0:h1]
            Kc = K[:, h0:h1]
            Vc = V[:, h0:h1]
            Ql = Q_land[:, h0:h1]
            Kl = K_land[:, h0:h1]

            with torch.amp.autocast("cuda", enabled=False):
                F1 = torch.matmul(Qc.float(), Kl.float().transpose(-2, -1).contiguous()) * scale
                F2 = torch.matmul(Ql.float(), Kl.float().transpose(-2, -1).contiguous()) * scale
                F3 = torch.matmul(Ql.float(), Kc.float().transpose(-2, -1).contiguous()) * scale

                # PR-FIX-01: pseudoinverse + Schulz without gradient tracking.
                # This eliminates the ~20x backward memory of pinv and stabilizes
                # the iteration numerically (rcond clamps near-singular values).
                with torch.no_grad():
                    F2_f32 = F2.detach().to(torch.float32)
                    F2_pinv = torch.linalg.pinv(F2_f32, rcond=1e-4)
                    I = (
                        torch.eye(
                            num_landmarks,
                            device=hidden_states.device,
                            dtype=torch.float32,
                        )
                        .unsqueeze(0)
                        .unsqueeze(0)
                    )
                    X = F2_pinv
                    for _ in range(min(self.pinv_iters, 4)):
                        X = torch.matmul(X, 2.0 * I - torch.matmul(F2_f32, X))
                    X = X.detach()

                    # PR-FIX-02: free pinv intermediates.
                    del F2_pinv, I

                attn_c = torch.matmul(F1, torch.matmul(X, F3))

                del F1, F2, F3, F2_f32, X

            # Causal mask
            is_causal = self.is_decoder or kwargs.get("is_causal", False)
            if is_causal:
                causal_mask = torch.triu(
                    torch.full((seq_len, seq_len), -1e4, device=attn_c.device),
                    diagonal=1,
                )
                attn_c = attn_c + causal_mask.unsqueeze(0).unsqueeze(0)

            # Padding mask
            if attention_mask is not None:
                attn_c = apply_attention_mask(attn_c, attention_mask)

            probs_c = F.softmax(attn_c, dim=-1).to(hidden_states.dtype)
            probs_c = torch.nan_to_num(probs_c, nan=0.0)
            probs_c = self.dropout(probs_c)

            # head_mask (per-head, so select chunk)
            if head_mask is not None:
                if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                    hm = head_mask[h0:h1].view(1, h1 - h0, 1, 1).to(probs_c.dtype)
                    probs_c = probs_c * hm
                elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                    hm = (
                        head_mask[:, h0:h1].view(head_mask.size(0), h1 - h0, 1, 1).to(probs_c.dtype)
                    )
                    probs_c = probs_c * hm
                elif head_mask.dim() == 4:
                    hm = head_mask[:, h0:h1].to(probs_c.dtype)
                    probs_c = probs_c * hm

            ctx_c = torch.matmul(probs_c, Vc)
            context_heads.append(ctx_c)

            if output_attentions and attn_for_return is None:
                # PR-FIX-05: return softmax probs, not logits.
                attn_for_return = probs_c

        # Concatenate chunks along head axis: (B, H_total, N, D)
        context = torch.cat(context_heads, dim=1)
        context = (
            context.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_len, self.all_head_size)
        )

        if self.is_nested_encoder:
            return context, (attn_for_return if output_attentions else None)

        out = self.out_proj(context)
        return out, (attn_for_return if output_attentions else None)

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class ProgressiveTokenPruning (config: Any | None = None,
d_model: int | None = None,
keep_ratio: float = 0.8,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
shared_scorer: SharedTokenScorer | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("progressive_token_pruning")
@ATTENTION_REGISTRY.register("ptp")
class ProgressiveTokenPruning(BaseAttention):
    """
    Progressive Token Pruning attention.

    Forward:
      1. Score each token (score_network or shared scorer).
      2. Force-keep the first token ([CLS]).
      3. Keep the top-k tokens (k = ceil(L · keep_ratio)).
      4. Compute attention over the reduced set of keys/values.
      5. Return output with the original (B, L, D) shape.

    KV-cache: supported for decoder inference.
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        d_model: int | None = None,
        keep_ratio: float = 0.8,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        old_attn: nn.Module | None = None,
        shared_scorer: SharedTokenScorer | None = None,
        **kwargs: Any,
    ) -> None:
        if rank is None:
            rank = 64
        super().__init__(config, rank)
        self.config = config
        comp_config = getattr(config, "compression_config", None)

        # PTP-FIX-12: dimension resolution.
        self.d_model = int(
            getattr(config, "hidden_size", None) or d_model or kwargs.get("hidden_size") or 768
        )
        self.dim = self.d_model
        self.num_heads = int(
            getattr(config, "num_attention_heads", kwargs.get("num_attention_heads", 12))
        )
        self.num_kv_heads = int(
            getattr(
                config,
                "num_key_value_heads",
                kwargs.get("num_key_value_heads", self.num_heads),
            )
        )
        self.head_dim = int(
            getattr(
                config,
                "head_dim",
                kwargs.get("head_dim", self.dim // max(1, self.num_heads)),
            )
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim
        self.keep_ratio = float(
            getattr(comp_config, "keep_ratio", kwargs.get("keep_ratio", keep_ratio))
        )

        is_dec = getattr(config, "is_decoder", kwargs.get("is_decoder", False))
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        self.lora_rank = int(
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = float(
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa = True

        # Projections
        self.q_lin = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_lin = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_lin = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_lin = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        # Aliases
        self.q_proj = self.q_lin
        self.k_proj = self.k_lin
        self.v_proj = self.v_lin
        self.out_proj = self.out_lin
        self.query = self.q_lin
        self.key = self.k_lin
        self.value = self.v_lin
        self.dense = self.out_lin  # PTP-FIX-10
        self.c_proj = self.out_lin  # PTP-FIX-10
        self.o_proj = self.out_lin  # PTP-FIX-10

        # PTP-FIX-04: optional shared scorer.
        if shared_scorer is not None:
            self.score_network = shared_scorer
            self._uses_shared_scorer = True
        else:
            self.score_network = SharedTokenScorer(self.dim, hidden=64)
            self._uses_shared_scorer = False

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(config, "attention_probs_dropout_prob", 0.1),
            )
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        src = (
            pretrained_attn
            if pretrained_attn is not None
            else (old_attn if old_attn is not None else kwargs.get("source_attn"))
        )
        if src is not None:
            self.copy_weights(src)

    # ------------------------------------------------------------------
    @property
    def gate(self):
        return None

    @property
    def random_proj(self):
        return None

    # ------------------------------------------------------------------
    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
        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[0] < c_w.shape[1])
                )
                if is_conv:
                    c_w = c_w.t().contiguous()
                split = 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, :],
                    c_b[:split] if c_b is not None else None,
                )
                extracted["k"] = (
                    c_w[split : 2 * split, :],
                    c_b[split : 2 * split] if c_b is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :, :],
                    c_b[2 * split :] 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. RoBERTa nested
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                mod = getattr(s, sub, None)
                if mod is not None and hasattr(mod, "weight"):
                    extracted[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
        else:
            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 key, names in candidates.items():
                for n in names:
                    if hasattr(src, n):
                        mod = getattr(src, n)
                        if mod is not None and hasattr(mod, "weight"):
                            extracted[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"ProgressiveTokenPruning copy_weights failed: missing "
                f"mandatory projections {missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_lin,
            "k": self.k_lin,
            "v": self.v_lin,
            "o": self.out_lin,
        }

        with torch.no_grad():
            for key, target_mod in target_map.items():
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    continue

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = float(torch.linalg.norm(target_mod.weight.float()).item())
                if fnorm < 1e-4:
                    logger.warning(
                        "PTP: '%s' weight has near-zero norm (%.2e).",
                        key,
                        fnorm,
                    )

        if self.freeze_base:
            for m in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
                m.freeze_base_weights()

    # ------------------------------------------------------------------
    # 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]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("PTP does not support cross-attention.")

        x = hidden_states
        B, L, D = x.shape

        valid_mask = get_token_valid_mask(attention_mask, B, L, x.device)

        # PTP-FIX-03: single assignment.
        scores = self.score_network(x)  # (B, L)
        # PTP-FIX-11: only force [CLS] if position 0 is valid.
        first_valid = valid_mask[:, 0].unsqueeze(-1)  # (B, 1)
        scores = torch.where(first_valid, scores, scores)
        # Force first token to top priority (only if valid).
        scores = scores.masked_fill(~valid_mask, -1e4)
        if first_valid.any():
            scores[:, 0] = torch.where(
                first_valid.squeeze(-1),
                torch.full_like(scores[:, 0], 1e4),
                scores[:, 0],
            )

        k_tokens = max(1, min(L, int(math.ceil(L * self.keep_ratio))))
        top_scores, top_indices = torch.topk(scores, k_tokens, dim=-1)
        top_indices, _ = torch.sort(top_indices, dim=-1)

        # Projections
        q = self.q_lin(x).view(B, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k_new = self.k_lin(x).view(B, L, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
        v_new = self.v_lin(x).view(B, L, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)

        # ---------------- KV-cache (PTP-FIX-01) ----------------
        past_k = past_v = None
        past_len = 0
        if past_key_value is not None and self.is_decoder is not False:
            # Tuple cache
            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]
                past_len = past_k.shape[-2]
            # HF DynamicCache
            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]
                past_len = past_k.shape[-2]

        # ---------------- RoPE with past_len offset (PTP-FIX-02) ----------------
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            if past_len > 0 and cos.shape[-2] >= past_len + L:
                cos = cos[..., past_len : past_len + L, :]
                sin = sin[..., past_len : past_len + L, :]
            q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin)

        # ---------------- Gather pruned K/V ----------------
        # Note: PTP prunes tokens ONLY for the current chunk (new tokens).
        # With KV-cache, past tokens are always used in full.
        top_idx_exp = (
            top_indices.unsqueeze(1).unsqueeze(-1).expand(-1, self.num_kv_heads, -1, self.head_dim)
        )
        k_pruned = torch.gather(k_new, 2, top_idx_exp)  # (B, H_kv, k_tokens, d)
        v_pruned = torch.gather(v_new, 2, top_idx_exp)

        if past_k is not None and past_v is not None:
            # Concatenate past (full) with current (pruned).
            k_all = torch.cat([past_k, k_pruned], dim=-2)
            v_all = torch.cat([past_v, v_pruned], dim=-2)
            # Keep track of original positions for causal masking.
            past_positions = torch.arange(past_len, device=x.device).view(1, 1, -1)
            past_positions = past_positions.expand(B, 1, -1)
            current_positions = (
                torch.gather(
                    top_indices,
                    1,
                    torch.arange(k_tokens, device=x.device).unsqueeze(0).expand(B, -1),
                )
                + past_len
            )  # positions in the full sequence
            k_positions = torch.cat(
                [past_positions, current_positions.unsqueeze(1)],
                dim=-1,
            )  # (B, 1, past_len + k_tokens)
        else:
            k_all = k_pruned
            v_all = v_pruned
            k_positions = top_indices.unsqueeze(1)  # (B, 1, k_tokens)

        # ---------------- GQA repeat ----------------
        if self.num_heads != self.num_kv_heads:
            k_all = repeat_kv(k_all, self.kv_repeat)
            v_all = repeat_kv(v_all, self.kv_repeat)

        # ---------------- Attention scores ----------------
        scale = 1.0 / math.sqrt(self.head_dim)
        attn_scores = (
            torch.matmul(q.float(), k_all.float().transpose(-2, -1).contiguous()) * scale
        )  # (B, H, L, past_len + k_tokens)

        # ---------------- Causal mask ----------------
        if self.is_decoder:
            q_pos = (past_len + torch.arange(L, device=x.device)).view(1, 1, L, 1)
            causal = k_positions.unsqueeze(2) > q_pos
            attn_scores = attn_scores.masked_fill(causal, -1e4)

        # ---------------- Valid-key mask for current chunk ----------------
        # Past keys are always valid (they were valid when cached).
        current_valid = torch.gather(valid_mask, 1, top_indices)  # (B, k_tokens)
        if past_k is not None:
            # Assume past keys are all valid (they were cached after masking).
            all_valid = torch.cat(
                [
                    torch.ones(B, past_len, device=x.device, dtype=torch.bool),
                    current_valid,
                ],
                dim=1,
            )  # (B, past_len + k_tokens)
        else:
            all_valid = current_valid
        attn_scores = attn_scores.masked_fill(~all_valid.unsqueeze(1).unsqueeze(2), -1e4)

        probs = F.softmax(attn_scores, dim=-1).to(x.dtype)
        probs = torch.nan_to_num(probs, nan=0.0)
        probs = self.dropout(probs)

        # head_mask handling
        if head_mask is not None:
            if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                probs = probs * head_mask.view(1, self.num_heads, 1, 1).to(probs.dtype)
            elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                probs = probs * head_mask.view(head_mask.size(0), self.num_heads, 1, 1).to(
                    probs.dtype
                )
            elif head_mask.dim() == 4:
                probs = probs * head_mask.to(probs.dtype)

        attn_out = (
            torch.matmul(probs, v_all)
            .permute(0, 2, 1, 3)
            .contiguous()
            .view(B, L, self.all_head_size)
        )

        # ---------------- Update KV-cache ----------------
        if past_key_value is not None:
            if hasattr(past_key_value, "update"):
                # HF DynamicCache: store the FULL k_new, v_new (not pruned),
                # so that future queries can access them.
                past_key_value.update(k_new, v_new, layer_idx)

        if self.is_nested_encoder:
            return attn_out, (probs if output_attentions else None)

        out = self.out_lin(attn_out)
        return out, (probs if output_attentions else None)

Progressive Token Pruning attention.

Forward -----= 1. Score each token (score_network or shared scorer). 2. Force-keep the first token ([CLS]). 3. Keep the top-k tokens (k = ceil(L · keep_ratio)). 4. Compute attention over the reduced set of keys/values. 5. Return output with the original (B, L, D) shape.

KV-cache: supported for decoder inference.

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.

Instance variables

prop gate

Expand source code

@property
def gate(self):
    return None

prop random_proj

Expand source code

@property
def random_proj(self):
    return None

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class ProjectionFactorizedAttention (config: Any | None = None,
old_attn: nn.Module | None = None,
rank: int | None = None,
mode: str = 'exact',
freeze_svd: bool = False,
lora_rank: int | None = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
use_decay_gating: bool = False,
decay_activation: str = 'softplus',
decay_channel_wise: bool = False,
use_fused_kernel: bool = False,
use_rms_norm: bool = False,
chunk_size: int = 64,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("projection_factorized_attn")
class ProjectionFactorizedAttention(SVDLinearAttention):
    r"""Projection Factorized SVD Linear Attention."""

    is_custom_attention: bool = True

Projection Factorized SVD Linear Attention.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• SVDLinearAttention
• BaseAttention
• torch.nn.modules.module.Module

Inherited members

• SVDLinearAttention:
  • copy_weights
  • extra_repr
  • forward
  • is_custom_attention
  • supports_bidirectional
  • supports_causal
  • supports_gqa
  • triton_kernel

class QCQAAttention (config: Any | None = None,
num_kv_heads: int | None = None,
router_temp: float = 1.0,
pretrained_attn: nn.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("qcqa")
@ATTENTION_REGISTRY.register("qcqa_attention")
@ATTENTION_REGISTRY.register("qcqa_attn")
class QCQAAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        num_kv_heads: int | None = None,
        router_temp: float = 1.0,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        rank = kwargs.pop("rank", 64)
        super().__init__(config, rank)
        self.config = config
        comp_config = kwargs.get("comp_config") or getattr(config, "compression_config", None)

        # hidden_size detection
        h_size = (
            kwargs.get("hidden_size")
            or getattr(config, "hidden_size", None)
            or getattr(config, "d_model", None)
        )
        if h_size is None and pretrained_attn is not None:
            for cand in ("q_proj", "q_lin", "query", "c_attn"):
                mod = getattr(pretrained_attn, cand, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    if w.ndim == 2:
                        if mod.__class__.__name__ == "Conv1D" or hasattr(mod, "nf"):
                            h_size = w.shape[0]
                        else:
                            h_size = w.shape[1]
                        break
        if h_size is None:
            raise ValueError("QCQAAttention: could not infer hidden_size.")
        self.dim = int(h_size)

        self.num_heads = int(
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or getattr(config, "num_attention_heads", None)
            or getattr(config, "num_heads", None)
            or 12
        )

        cfg_kv = getattr(config, "num_key_value_heads", getattr(config, "num_kv_heads", None))
        if num_kv_heads is not None:
            self.num_kv_heads = int(num_kv_heads)
        elif cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        else:
            self.num_kv_heads = int(kwargs.get("num_kv_heads", 4))
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))

        self.head_dim = int(
            kwargs.get("head_dim")
            or getattr(config, "head_dim", None)
            or getattr(config, "d_kv", None)
            or (self.dim // max(1, self.num_heads))
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        self.router_temp = float(
            _cfg_get(comp_config, "router_temp", kwargs.get("router_temp", router_temp))
        )
        if self.router_temp <= 0:
            self.router_temp = 1.0

        self.is_decoder = _detect_decoder(config, is_decoder, pretrained_attn)
        self.is_nested_encoder = False

        self.lora_rank = int(
            _cfg_get(
                config,
                "lora_rank",
                _cfg_get(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
        )
        self.lora_alpha = float(
            _cfg_get(
                config,
                "lora_alpha",
                _cfg_get(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        )
        self.freeze_base = self.lora_rank > 0
        self.supports_gqa = True

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and pretrained_attn is not None:
            p = next(pretrained_attn.parameters(), None)
            if p is not None:
                target_dtype = p.dtype

        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )

        # Aliases
        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj
        self.q_lin = self.q_proj
        self.k_lin = self.k_proj
        self.v_lin = self.v_proj
        self.out_lin = self.out_proj
        self.dense = self.out_proj
        self.o_proj = self.out_proj
        self.c_proj = self.out_proj

        # QC-FIX-20: router maps q's last dim → num_kv_heads
        self.router = nn.Linear(self.head_dim, self.num_kv_heads, bias=False)
        nn.init.normal_(self.router.weight, std=0.02)

        dropout_p = _cfg_get(
            config, "dropout", _cfg_get(config, "attention_probs_dropout_prob", 0.1)
        )
        self.dropout = nn.Dropout(dropout_p)

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if source_attn is not None:
            self.copy_weights(source_attn)

        if target_dtype is not None:
            self.to(target_dtype)

    # ------------------------------------------------------------------
    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",
            )
            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),
        }

        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()
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        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])
                    ):
                        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,
                    )

        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,
                )

        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"QCQAAttention.copy_weights: missing mandatory projections "
                f"{missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    else:
                        raise RuntimeError(f"QCQAAttention: missing weight for {proj_key}.")

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                    and w.shape[1] == target_mod.weight.shape[1]
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                norm = float(torch.linalg.norm(target_mod.weight.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "QCQAAttention: weight for '%s' has near-zero Frobenius norm (%.2e).",
                        proj_key,
                        norm,
                    )

        if self.freeze_base:
            for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
                mod.freeze_base_weights()

    # ------------------------------------------------------------------
    # QC-FIX-16: robust valid-mask helper (2D / 3D / 4D)
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 3:
            # (B, N_q, N_kv) or (B, 1, N_kv)
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            # If query axis has size > 1, aggregate (any valid).
            if v.shape[1] > 1:
                v = v.any(dim=1)
            else:
                v = v.squeeze(1)
            return v  # (B, N_kv)
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)  # (B, N_kv)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # 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, L, _ = 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]

        # --- Projections ---
        q = self.q_proj(hidden_states).view(B, L, self.num_heads, self.head_dim)
        k_new = self.k_proj(kv_input).view(B, N_kv_new, self.num_kv_heads, self.head_dim)
        v_new = self.v_proj(kv_input).view(B, N_kv_new, self.num_kv_heads, self.head_dim)

        # --- RoPE (self-attention only) ---
        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:
                    idx = min(layer_idx, len(past_key_value.key_cache) - 1)
                    past_len_rope = past_key_value.key_cache[idx].shape[-2]
            # QC-FIX-19: cos/sin must cover both q and k new tokens.
            cos_slice = cos[..., past_len_rope : past_len_rope + N_kv_new, :]
            sin_slice = sin[..., past_len_rope : past_len_rope + N_kv_new, :]

            # q uses the first L of these positions; k uses the last N_kv_new (== same if no cross)
            q_t = q.permute(0, 2, 1, 3).contiguous()  # (B, H_q, L, d)
            k_t = k_new.permute(0, 2, 1, 3).contiguous()  # (B, H_kv, N_kv_new, d)

            # Use slices matching L for q and N_kv_new for k.
            cos_q = cos_slice[..., :L, :]
            sin_q = sin_slice[..., :L, :]
            cos_k = cos_slice[..., :N_kv_new, :]
            sin_k = sin_slice[..., :N_kv_new, :]
            q_rot, _ = apply_rotary_pos_emb(q_t, q_t, cos_q, sin_q)
            _, k_rot = apply_rotary_pos_emb(k_t, k_t, cos_k, sin_k)
            q = q_rot.permute(0, 2, 1, 3).contiguous()
            k_new = k_rot.permute(0, 2, 1, 3).contiguous()

        # --- 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=1)  # (B, N_kv, H_kv, d)
            V_cat = torch.cat([past_v, v_new], dim=1)
        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[1]
        past_len = N_kv - N_kv_new if (self.is_decoder and not is_cross) else 0

        # ============================================================
        # QC-FIX-13: Dynamic routing → head-level KV synthesis → MHA
        # ============================================================
        # Router: per-token distribution over KV heads
        # q: (B, L, H_q, d) → router → (B, L, H_q, H_kv)
        route_logits = self.router(q) / max(self.router_temp, 1e-4)
        route_weights = F.softmax(route_logits.float(), dim=-1).to(q.dtype)  # (B, L, H_q, H_kv)

        # Aggregate over query positions to get head-level routing.
        # This keeps the routing input-dependent (from q) while remaining
        # computationally tractable and semantically meaningful.
        w_head = route_weights.mean(dim=1)  # (B, H_q, H_kv)
        w_head = w_head / w_head.sum(dim=-1, keepdim=True).clamp(min=1e-9)

        # Mix KV heads: k_cat[b, h_q, j, d] = Σ_{h_kv} w_head[b, h_q, h_kv] · k[b, j, h_kv, d]
        K_perm = K_cat.permute(0, 2, 1, 3)  # (B, H_kv, N_kv, d)
        V_perm = V_cat.permute(0, 2, 1, 3)
        # einsum: w_head[b,h_q,h_kv] · K_perm[b,h_kv,j,d] → (B, H_q, N_kv, d)
        K_mixed = torch.einsum("bqk,bkjd->bqjd", w_head, K_perm)  # (B, H_q, N_kv, d)
        V_mixed = torch.einsum("bqk,bkjd->bqjd", w_head, V_perm)

        # Prepare q for attention
        q_heads = q.permute(0, 2, 1, 3).contiguous()  # (B, H_q, L, d)

        scale = 1.0 / math.sqrt(self.head_dim)

        # Causal mask (decoder only)
        causal_mask = None
        if self.is_decoder and not is_cross:
            if past_len > 0:
                abs_q = (past_len + torch.arange(L, device=device)).unsqueeze(1)
                abs_k = torch.arange(N_kv, device=device).unsqueeze(0)
                causal = (abs_k <= abs_q).unsqueeze(0).unsqueeze(0)  # (1, 1, L, N_kv)
            else:
                causal = (
                    torch.tril(torch.ones(L, N_kv, device=device, dtype=torch.bool))
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
            causal_mask = causal

        # Padding mask
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

        # --- Fast SDPA path (QC-FIX-18) ---
        # Build additive bias for SDPA
        attn_bias_sdpa = None
        if not output_attentions and head_mask is None:
            bias_parts = []
            if causal_mask is not None:
                causal_bias = torch.zeros(1, 1, L, N_kv, device=device, dtype=torch.float32)
                causal_bias.masked_fill_(~causal_mask, -1e4)
                bias_parts.append(causal_bias)
            if token_valid is not None:
                pad_bias = torch.zeros(1, 1, 1, N_kv, device=device, dtype=torch.float32)
                pad_bias.masked_fill_(~token_valid[:, None, None, :].bool(), -1e4)
                bias_parts.append(pad_bias)
            if bias_parts:
                combined = sum(bias_parts)
                attn_bias_sdpa = combined.to(K_mixed.dtype)

        if attn_bias_sdpa is not None or (causal_mask is None and token_valid is None):
            try:
                q_sdp = q_heads.to(torch.float16) if q_heads.dtype == torch.float32 else q_heads
                k_sdp = K_mixed.to(torch.float16) if K_mixed.dtype == torch.float32 else K_mixed
                v_sdp = V_mixed.to(torch.float16) if V_mixed.dtype == torch.float32 else V_mixed

                out_sdp = F.scaled_dot_product_attention(
                    q_sdp,
                    k_sdp,
                    v_sdp,
                    attn_mask=attn_bias_sdpa,
                    is_causal=False,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                context = (
                    out_sdp.to(q.dtype)
                    .permute(0, 2, 1, 3)
                    .contiguous()
                    .view(B, L, self.all_head_size)
                )
                context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

                if self.is_nested_encoder:
                    return context, None
                out = self.out_proj(context)
                out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
                return out, None
            except Exception as e:
                logger.debug("QCQA SDPA path failed: %s; falling back to manual.", e)

        # --- Manual path (needed when output_attentions or head_mask) ---
        scores = (
            torch.matmul(
                q_heads.float(),
                K_mixed.float().transpose(-2, -1).contiguous(),
            )
            * scale
        )  # (B, H_q, L, N_kv)

        if causal_mask is not None:
            scores = scores.masked_fill(~causal_mask, -1e4)

        if token_valid is not None:
            scores = scores.masked_fill(~token_valid[:, None, None, :].bool(), -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_probs = F.softmax(scores, dim=-1).to(q.dtype)
        attn_probs = torch.nan_to_num(attn_probs, nan=0.0)
        attn_probs = self.dropout(attn_probs)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)
            attn_probs = attn_probs * head_mask.to(attn_probs.dtype)

        # Context: (B, H_q, L, N_kv) @ (B, H_q, N_kv, d) → (B, H_q, L, d)
        context = torch.matmul(attn_probs, V_mixed)
        context = context.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            return context, (attn_probs if output_attentions else None)

        out = self.out_proj(context)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
        return out, (attn_probs if output_attentions else None)

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class RALACompressorAttention (config: Any | None = None,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("rala")
@ATTENTION_REGISTRY.register("rala_attention")
@ATTENTION_REGISTRY.register("rala_linear_attention")
class RALALinearAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank or 64)
        except TypeError:
            super().__init__(config)
        self.config = config
        comp_config = getattr(config, "compression_config", None)

        self.hidden_size = getattr(config, "hidden_size", kwargs.get("hidden_size", 768))
        if (
            (config is None or not hasattr(config, "hidden_size"))
            and "hidden_size" not in kwargs
            and pretrained_attn is not None
        ):
            for cand in ["q_lin", "query", "q_proj", "c_attn"]:
                mod = getattr(pretrained_attn, cand, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    self.hidden_size = w.shape[1] if w.ndim == 2 else self.hidden_size
                    break

        self.dim = self.hidden_size
        self.num_heads = getattr(
            config, "num_attention_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads = getattr(
            config,
            "num_key_value_heads",
            kwargs.get("num_key_value_heads", self.num_heads),
        )
        self.head_dim = getattr(
            config,
            "head_dim",
            kwargs.get("head_dim", self.dim // self.num_heads),
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        self.rank_r = getattr(
            comp_config,
            "rala_rank",
            getattr(config, "rala_rank", kwargs.get("rala_rank", 8)),
        )

        # RALA-FIX-01: rename trainable params with `lora_` prefix
        self.lora_U = nn.Parameter(torch.randn(self.num_heads, self.head_dim, self.rank_r) * 0.01)
        self.lora_V = nn.Parameter(torch.randn(self.num_heads, self.head_dim, self.rank_r) * 0.01)
        self.lora_rank_alpha = nn.Parameter(torch.tensor(0.0))

        # Decoder detection
        _mt = str(getattr(config, "model_type", "") or "").lower()
        _decoder_prefixes = ("gpt", "qwen", "llama", "mistral", "gemma", "phi", "opt")
        self.is_decoder = (
            bool(getattr(config, "is_decoder", False))
            or bool(getattr(config, "causal", False))
            or any(_mt.startswith(p) for p in _decoder_prefixes)
            or (pretrained_attn is not None and hasattr(pretrained_attn, "c_attn"))
            or bool(kwargs.get("is_causal", False))
        )
        if "is_decoder" in kwargs and kwargs["is_decoder"] is not None:
            self.is_decoder = bool(kwargs["is_decoder"])
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa = True

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        target_dtype = kwargs.get("dtype")
        if target_dtype is None and source_attn is not None:
            p = next(source_attn.parameters(), None)
            if p is not None:
                target_dtype = p.dtype

        # Projections
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )

        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        if source_attn is not None:
            self.copy_weights(source_attn)
            try:
                src_device = next(source_attn.parameters()).device
                self.to(src_device)
            except Exception:
                pass

        if target_dtype is not None:
            self.to(target_dtype)

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    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
        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()
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    c_w = c_w.t()
                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()
                    if (
                        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])
                    ):
                        o_w = o_w.t()
                    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. 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
        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:
            logger.warning(
                "RALALinearAttention.copy_weights: missing %s from %s",
                missing_after,
                src.__class__.__name__,
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    continue

                src_norm = float(torch.linalg.norm(w.float()).item())
                if src_norm < 1e-4:
                    logger.warning(
                        "RALALinearAttention: source weight for %s has near-zero norm %.6f",
                        proj_key,
                        src_norm,
                    )

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                    and w.shape[1] == target_mod.weight.shape[1]
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = float(torch.linalg.norm(target_mod.weight.float()).item())
                if fnorm < 1e-4:
                    logger.warning(
                        "RALALinearAttention: %s has norm %.6f",
                        proj_key,
                        fnorm,
                    )

        if self.freeze_base:
            for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
                mod.freeze_base_weights()
            # RALA-FIX-01: U, V, alpha trainable via lora_ prefix
            self.lora_U.requires_grad = True
            self.lora_V.requires_grad = True
            self.lora_rank_alpha.requires_grad = True

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, seq_len, _ = hidden_states.shape

        q = (
            self.q_proj(hidden_states)
            .view(batch_size, seq_len, self.num_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        k = (
            self.k_proj(hidden_states)
            .view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        v = (
            self.v_proj(hidden_states)
            .view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )

        # RoPE
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # GQA
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = repeat_kv(k, repeats)
            v = repeat_kv(v, repeats)

        q_prime = F.softplus(q)

        valid_mask = get_token_valid_mask(attention_mask, batch_size, seq_len, hidden_states.device)
        mask_4d = valid_mask.unsqueeze(1).unsqueeze(-1).to(k.dtype)
        k_prime = F.softplus(k) * mask_4d
        v_masked = v * mask_4d

        is_causal = self.is_decoder or kwargs.get("is_causal", False)

        if is_causal:
            # State-based linear scan
            S = torch.zeros(
                batch_size,
                self.num_heads,
                self.head_dim,
                self.head_dim,
                dtype=q.dtype,
                device=hidden_states.device,
            )
            z = torch.zeros(
                batch_size,
                self.num_heads,
                self.head_dim,
                dtype=q.dtype,
                device=hidden_states.device,
            )

            num_list = []
            denom_list = []

            for i in range(seq_len):
                q_i = q_prime[:, :, i, :]
                k_i = k_prime[:, :, i, :]
                v_i = v_masked[:, :, i, :]

                S = S + torch.einsum("bhd,bhe->bhde", k_i, v_i)
                z = z + k_i

                num_i = torch.einsum("bhd,bhde->bhe", q_i, S)
                denom_i = torch.einsum("bhd,bhd->bh", q_i, z).unsqueeze(-1)

                num_list.append(num_i)
                denom_list.append(denom_i)

            num = torch.stack(num_list, dim=2)
            denom = torch.stack(denom_list, dim=2).clamp(min=1e-4)
            lin_out = torch.nan_to_num(num / denom)

            # RALA-FIX-01: use lora_U, lora_V (trainable)
            rank_mat = torch.einsum("hdr,her->hde", self.lora_U, self.lora_V)
            aug_out = torch.nan_to_num(torch.einsum("bhnd,hde->bhne", q_prime, rank_mat) / denom)
        else:
            kv = torch.einsum("bhnd,bhne->bhde", k_prime, v_masked)
            z = k_prime.sum(dim=2)
            denom = torch.einsum("bhnd,bhd->bhn", q_prime, z).unsqueeze(-1).clamp(min=1e-4)
            lin_out = torch.nan_to_num(torch.einsum("bhnd,bhde->bhne", q_prime, kv) / denom)

            rank_mat = torch.einsum("hdr,her->hde", self.lora_U, self.lora_V)
            aug_out = torch.nan_to_num(torch.einsum("bhnd,hde->bhne", q_prime, rank_mat) / denom)

        # RALA-FIX-04: removed dead scale_factor=1.0
        # RALA-FIX-01: use lora_rank_alpha (trainable)
        out = (lin_out + self.lora_rank_alpha * aug_out) * mask_4d

        # head_mask
        if head_mask is not None:
            if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                out = out * head_mask.view(1, self.num_heads, 1, 1).to(out.dtype)
            elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                out = out * head_mask.view(head_mask.size(0), self.num_heads, 1, 1).to(out.dtype)
            elif head_mask.dim() == 4:
                out = out * head_mask.to(out.dtype)

        out = out.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_len, self.all_head_size)

        if self.is_nested_encoder:
            return out, None

        output = self.out_proj(out)
        return output, None

Every 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.

Methods

def copy_weights(self,
src: Any) ‑> None

Expand source code

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
    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()
            if (
                c_attn_mod.__class__.__name__ == "Conv1D"
                or base_c.__class__.__name__ == "Conv1D"
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                c_w = c_w.t()
            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()
                if (
                    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])
                ):
                    o_w = o_w.t()
                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. 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
    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:
        logger.warning(
            "RALALinearAttention.copy_weights: missing %s from %s",
            missing_after,
            src.__class__.__name__,
        )

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                    continue
                continue

            src_norm = float(torch.linalg.norm(w.float()).item())
            if src_norm < 1e-4:
                logger.warning(
                    "RALALinearAttention: source weight for %s has near-zero norm %.6f",
                    proj_key,
                    src_norm,
                )

            if w.shape == target_mod.weight.shape:
                target_mod.weight.copy_(w)
            elif w.t().shape == target_mod.weight.shape:
                target_mod.weight.copy_(w.t())
            elif (
                w.shape[0] < target_mod.weight.shape[0]
                and target_mod.weight.shape[0] % w.shape[0] == 0
                and w.shape[1] == target_mod.weight.shape[1]
            ):
                repeats = target_mod.weight.shape[0] // w.shape[0]
                target_mod.weight.copy_(w.repeat(repeats, 1))
            else:
                min_0 = min(target_mod.weight.shape[0], w.shape[0])
                min_1 = min(target_mod.weight.shape[1], w.shape[1])
                target_mod.weight.zero_()
                target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

            if b is not None and target_mod.bias is not None:
                if b.shape == target_mod.bias.shape:
                    target_mod.bias.copy_(b)
                elif (
                    b.shape[0] < target_mod.bias.shape[0]
                    and target_mod.bias.shape[0] % b.shape[0] == 0
                ):
                    repeats = target_mod.bias.shape[0] // b.shape[0]
                    target_mod.bias.copy_(b.repeat(repeats))
                else:
                    min_b = min(target_mod.bias.shape[0], b.shape[0])
                    target_mod.bias.zero_()
                    target_mod.bias[:min_b].copy_(b[:min_b])
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

            fnorm = float(torch.linalg.norm(target_mod.weight.float()).item())
            if fnorm < 1e-4:
                logger.warning(
                    "RALALinearAttention: %s has norm %.6f",
                    proj_key,
                    fnorm,
                )

    if self.freeze_base:
        for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
            mod.freeze_base_weights()
        # RALA-FIX-01: U, V, alpha trainable via lora_ prefix
        self.lora_U.requires_grad = True
        self.lora_V.requires_grad = True
        self.lora_rank_alpha.requires_grad = True

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    batch_size, seq_len, _ = hidden_states.shape

    q = (
        self.q_proj(hidden_states)
        .view(batch_size, seq_len, self.num_heads, self.head_dim)
        .permute(0, 2, 1, 3)
    )
    k = (
        self.k_proj(hidden_states)
        .view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        .permute(0, 2, 1, 3)
    )
    v = (
        self.v_proj(hidden_states)
        .view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        .permute(0, 2, 1, 3)
    )

    # RoPE
    if position_embeddings is None:
        position_embeddings = kwargs.get("position_embeddings")
    if position_embeddings is not None:
        cos, sin = position_embeddings
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    # GQA
    if self.num_heads != self.num_kv_heads:
        repeats = self.num_heads // self.num_kv_heads
        k = repeat_kv(k, repeats)
        v = repeat_kv(v, repeats)

    q_prime = F.softplus(q)

    valid_mask = get_token_valid_mask(attention_mask, batch_size, seq_len, hidden_states.device)
    mask_4d = valid_mask.unsqueeze(1).unsqueeze(-1).to(k.dtype)
    k_prime = F.softplus(k) * mask_4d
    v_masked = v * mask_4d

    is_causal = self.is_decoder or kwargs.get("is_causal", False)

    if is_causal:
        # State-based linear scan
        S = torch.zeros(
            batch_size,
            self.num_heads,
            self.head_dim,
            self.head_dim,
            dtype=q.dtype,
            device=hidden_states.device,
        )
        z = torch.zeros(
            batch_size,
            self.num_heads,
            self.head_dim,
            dtype=q.dtype,
            device=hidden_states.device,
        )

        num_list = []
        denom_list = []

        for i in range(seq_len):
            q_i = q_prime[:, :, i, :]
            k_i = k_prime[:, :, i, :]
            v_i = v_masked[:, :, i, :]

            S = S + torch.einsum("bhd,bhe->bhde", k_i, v_i)
            z = z + k_i

            num_i = torch.einsum("bhd,bhde->bhe", q_i, S)
            denom_i = torch.einsum("bhd,bhd->bh", q_i, z).unsqueeze(-1)

            num_list.append(num_i)
            denom_list.append(denom_i)

        num = torch.stack(num_list, dim=2)
        denom = torch.stack(denom_list, dim=2).clamp(min=1e-4)
        lin_out = torch.nan_to_num(num / denom)

        # RALA-FIX-01: use lora_U, lora_V (trainable)
        rank_mat = torch.einsum("hdr,her->hde", self.lora_U, self.lora_V)
        aug_out = torch.nan_to_num(torch.einsum("bhnd,hde->bhne", q_prime, rank_mat) / denom)
    else:
        kv = torch.einsum("bhnd,bhne->bhde", k_prime, v_masked)
        z = k_prime.sum(dim=2)
        denom = torch.einsum("bhnd,bhd->bhn", q_prime, z).unsqueeze(-1).clamp(min=1e-4)
        lin_out = torch.nan_to_num(torch.einsum("bhnd,bhde->bhne", q_prime, kv) / denom)

        rank_mat = torch.einsum("hdr,her->hde", self.lora_U, self.lora_V)
        aug_out = torch.nan_to_num(torch.einsum("bhnd,hde->bhne", q_prime, rank_mat) / denom)

    # RALA-FIX-04: removed dead scale_factor=1.0
    # RALA-FIX-01: use lora_rank_alpha (trainable)
    out = (lin_out + self.lora_rank_alpha * aug_out) * mask_4d

    # head_mask
    if head_mask is not None:
        if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
            out = out * head_mask.view(1, self.num_heads, 1, 1).to(out.dtype)
        elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
            out = out * head_mask.view(head_mask.size(0), self.num_heads, 1, 1).to(out.dtype)
        elif head_mask.dim() == 4:
            out = out * head_mask.to(out.dtype)

    out = out.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_len, self.all_head_size)

    if self.is_nested_encoder:
        return out, None

    output = self.out_proj(out)
    return output, None

Must return (output, attention_weights_or_None).

class RALALinearAttention (config: Any | None = None,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("rala")
@ATTENTION_REGISTRY.register("rala_attention")
@ATTENTION_REGISTRY.register("rala_linear_attention")
class RALALinearAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank or 64)
        except TypeError:
            super().__init__(config)
        self.config = config
        comp_config = getattr(config, "compression_config", None)

        self.hidden_size = getattr(config, "hidden_size", kwargs.get("hidden_size", 768))
        if (
            (config is None or not hasattr(config, "hidden_size"))
            and "hidden_size" not in kwargs
            and pretrained_attn is not None
        ):
            for cand in ["q_lin", "query", "q_proj", "c_attn"]:
                mod = getattr(pretrained_attn, cand, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    self.hidden_size = w.shape[1] if w.ndim == 2 else self.hidden_size
                    break

        self.dim = self.hidden_size
        self.num_heads = getattr(
            config, "num_attention_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads = getattr(
            config,
            "num_key_value_heads",
            kwargs.get("num_key_value_heads", self.num_heads),
        )
        self.head_dim = getattr(
            config,
            "head_dim",
            kwargs.get("head_dim", self.dim // self.num_heads),
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        self.rank_r = getattr(
            comp_config,
            "rala_rank",
            getattr(config, "rala_rank", kwargs.get("rala_rank", 8)),
        )

        # RALA-FIX-01: rename trainable params with `lora_` prefix
        self.lora_U = nn.Parameter(torch.randn(self.num_heads, self.head_dim, self.rank_r) * 0.01)
        self.lora_V = nn.Parameter(torch.randn(self.num_heads, self.head_dim, self.rank_r) * 0.01)
        self.lora_rank_alpha = nn.Parameter(torch.tensor(0.0))

        # Decoder detection
        _mt = str(getattr(config, "model_type", "") or "").lower()
        _decoder_prefixes = ("gpt", "qwen", "llama", "mistral", "gemma", "phi", "opt")
        self.is_decoder = (
            bool(getattr(config, "is_decoder", False))
            or bool(getattr(config, "causal", False))
            or any(_mt.startswith(p) for p in _decoder_prefixes)
            or (pretrained_attn is not None and hasattr(pretrained_attn, "c_attn"))
            or bool(kwargs.get("is_causal", False))
        )
        if "is_decoder" in kwargs and kwargs["is_decoder"] is not None:
            self.is_decoder = bool(kwargs["is_decoder"])
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa = True

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        target_dtype = kwargs.get("dtype")
        if target_dtype is None and source_attn is not None:
            p = next(source_attn.parameters(), None)
            if p is not None:
                target_dtype = p.dtype

        # Projections
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )

        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        if source_attn is not None:
            self.copy_weights(source_attn)
            try:
                src_device = next(source_attn.parameters()).device
                self.to(src_device)
            except Exception:
                pass

        if target_dtype is not None:
            self.to(target_dtype)

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    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
        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()
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    c_w = c_w.t()
                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()
                    if (
                        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])
                    ):
                        o_w = o_w.t()
                    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. 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
        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:
            logger.warning(
                "RALALinearAttention.copy_weights: missing %s from %s",
                missing_after,
                src.__class__.__name__,
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    continue

                src_norm = float(torch.linalg.norm(w.float()).item())
                if src_norm < 1e-4:
                    logger.warning(
                        "RALALinearAttention: source weight for %s has near-zero norm %.6f",
                        proj_key,
                        src_norm,
                    )

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                    and w.shape[1] == target_mod.weight.shape[1]
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = float(torch.linalg.norm(target_mod.weight.float()).item())
                if fnorm < 1e-4:
                    logger.warning(
                        "RALALinearAttention: %s has norm %.6f",
                        proj_key,
                        fnorm,
                    )

        if self.freeze_base:
            for mod in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
                mod.freeze_base_weights()
            # RALA-FIX-01: U, V, alpha trainable via lora_ prefix
            self.lora_U.requires_grad = True
            self.lora_V.requires_grad = True
            self.lora_rank_alpha.requires_grad = True

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, seq_len, _ = hidden_states.shape

        q = (
            self.q_proj(hidden_states)
            .view(batch_size, seq_len, self.num_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        k = (
            self.k_proj(hidden_states)
            .view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        v = (
            self.v_proj(hidden_states)
            .view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )

        # RoPE
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # GQA
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = repeat_kv(k, repeats)
            v = repeat_kv(v, repeats)

        q_prime = F.softplus(q)

        valid_mask = get_token_valid_mask(attention_mask, batch_size, seq_len, hidden_states.device)
        mask_4d = valid_mask.unsqueeze(1).unsqueeze(-1).to(k.dtype)
        k_prime = F.softplus(k) * mask_4d
        v_masked = v * mask_4d

        is_causal = self.is_decoder or kwargs.get("is_causal", False)

        if is_causal:
            # State-based linear scan
            S = torch.zeros(
                batch_size,
                self.num_heads,
                self.head_dim,
                self.head_dim,
                dtype=q.dtype,
                device=hidden_states.device,
            )
            z = torch.zeros(
                batch_size,
                self.num_heads,
                self.head_dim,
                dtype=q.dtype,
                device=hidden_states.device,
            )

            num_list = []
            denom_list = []

            for i in range(seq_len):
                q_i = q_prime[:, :, i, :]
                k_i = k_prime[:, :, i, :]
                v_i = v_masked[:, :, i, :]

                S = S + torch.einsum("bhd,bhe->bhde", k_i, v_i)
                z = z + k_i

                num_i = torch.einsum("bhd,bhde->bhe", q_i, S)
                denom_i = torch.einsum("bhd,bhd->bh", q_i, z).unsqueeze(-1)

                num_list.append(num_i)
                denom_list.append(denom_i)

            num = torch.stack(num_list, dim=2)
            denom = torch.stack(denom_list, dim=2).clamp(min=1e-4)
            lin_out = torch.nan_to_num(num / denom)

            # RALA-FIX-01: use lora_U, lora_V (trainable)
            rank_mat = torch.einsum("hdr,her->hde", self.lora_U, self.lora_V)
            aug_out = torch.nan_to_num(torch.einsum("bhnd,hde->bhne", q_prime, rank_mat) / denom)
        else:
            kv = torch.einsum("bhnd,bhne->bhde", k_prime, v_masked)
            z = k_prime.sum(dim=2)
            denom = torch.einsum("bhnd,bhd->bhn", q_prime, z).unsqueeze(-1).clamp(min=1e-4)
            lin_out = torch.nan_to_num(torch.einsum("bhnd,bhde->bhne", q_prime, kv) / denom)

            rank_mat = torch.einsum("hdr,her->hde", self.lora_U, self.lora_V)
            aug_out = torch.nan_to_num(torch.einsum("bhnd,hde->bhne", q_prime, rank_mat) / denom)

        # RALA-FIX-04: removed dead scale_factor=1.0
        # RALA-FIX-01: use lora_rank_alpha (trainable)
        out = (lin_out + self.lora_rank_alpha * aug_out) * mask_4d

        # head_mask
        if head_mask is not None:
            if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                out = out * head_mask.view(1, self.num_heads, 1, 1).to(out.dtype)
            elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                out = out * head_mask.view(head_mask.size(0), self.num_heads, 1, 1).to(out.dtype)
            elif head_mask.dim() == 4:
                out = out * head_mask.to(out.dtype)

        out = out.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_len, self.all_head_size)

        if self.is_nested_encoder:
            return out, None

        output = self.out_proj(out)
        return output, None

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class REAPMoE (d_model: int = 768,
n_experts: int = 2,
top_k: int = 1,
noise_std: float = 0.05,
use_svd: bool = True,
config: Any | None = None,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
**kwargs: Any)

Expand source code

class REAPMoE(nn.Module):
    def __init__(
        self,
        d_model: int = 768,
        n_experts: int = 2,
        top_k: int = 1,
        noise_std: float = 0.05,
        use_svd: bool = True,
        config: Any | None = None,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__()
        factory_kwargs = {"device": device, "dtype": dtype}
        if config is not None:
            d_model = getattr(config, "hidden_size", getattr(config, "d_model", d_model))
            n_experts = getattr(config, "num_experts", getattr(config, "n_experts", n_experts))
            top_k = getattr(config, "top_k", top_k)

        self.d_model = d_model
        self.n_experts = n_experts
        self.top_k = min(top_k, n_experts)
        self.noise_std = noise_std

        self.register_buffer("_aux_loss", torch.tensor(0.0), persistent=False)

        self.router_ln = nn.LayerNorm(self.d_model, **factory_kwargs)
        self.router = nn.Linear(d_model, n_experts, bias=False, **factory_kwargs)
        nn.init.uniform_(self.router.weight, -0.01, 0.01)

        if use_svd:

            def expert_layer():
                return nn.Sequential(
                    SVDLinear(d_model, d_model, rank=max(1, d_model // 8), **factory_kwargs),
                    nn.GELU(),
                    SVDLinear(d_model, d_model, rank=max(1, d_model // 8), **factory_kwargs),
                )

        else:

            def expert_layer():
                return nn.Sequential(
                    nn.Linear(d_model, d_model, **factory_kwargs),
                    nn.GELU(),
                    nn.Linear(d_model, d_model, **factory_kwargs),
                )

        self.experts = nn.ModuleList([expert_layer() for _ in range(n_experts)])

        for exp in self.experts:
            final_layer = exp[-1]
            if isinstance(final_layer, SVDLinear):
                nn.init.zeros_(final_layer.U)
                nn.init.zeros_(final_layer.S)
                nn.init.zeros_(final_layer.bias)
            elif isinstance(final_layer, nn.Linear):
                nn.init.zeros_(final_layer.weight)
                if final_layer.bias is not None:
                    nn.init.zeros_(final_layer.bias)

    @property
    def aux_loss(self) -> torch.Tensor:
        return self._aux_loss

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        B, L, D = x.shape
        target_dtype = x.dtype

        r_dtype = (
            self.router_ln.weight.dtype
            if hasattr(self.router_ln, "weight") and self.router_ln.weight is not None
            else target_dtype
        )
        normed_x = self.router_ln(x.to(r_dtype)).to(target_dtype)
        rw_dtype = (
            self.router.weight.dtype
            if hasattr(self.router, "weight") and self.router.weight is not None
            else target_dtype
        )
        logits = self.router(normed_x.to(rw_dtype)).to(target_dtype)
        if self.training and self.noise_std > 0:
            noise = torch.randn_like(logits) * self.noise_std
            scores = logits + noise
        else:
            scores = logits

        full_probs = F.softmax(scores.float(), dim=-1).to(target_dtype)
        topk_weights, topk_indices = torch.topk(full_probs, self.top_k, dim=-1)
        topk_weights = (topk_weights / topk_weights.sum(dim=-1, keepdim=True).clamp(min=1e-9)).to(
            target_dtype
        )

        router_probs = F.softmax(logits.float(), dim=-1).view(-1, self.n_experts)
        expert_mask = F.one_hot(topk_indices, num_classes=self.n_experts).float()
        expert_assignment = expert_mask.sum(dim=2).clamp(max=1.0)
        density = expert_assignment.view(-1, self.n_experts).mean(dim=0)
        prob_mean = router_probs.mean(dim=0)
        aux_loss = self.n_experts * torch.sum(density * prob_mean)
        self._aux_loss = aux_loss.detach()

        experts_stack = torch.stack([exp(x).to(target_dtype) for exp in self.experts], dim=2)
        topk_expanded = topk_indices.unsqueeze(-1).expand(-1, -1, -1, D)
        gathered = torch.gather(experts_stack, 2, topk_expanded)
        out = torch.sum(gathered * topk_weights.unsqueeze(-1), dim=2).to(target_dtype)

        return out, aux_loss

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Instance variables

prop aux_loss : torch.Tensor

Expand source code

@property
def aux_loss(self) -> torch.Tensor:
    return self._aux_loss

Methods

def forward(self, x: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    B, L, D = x.shape
    target_dtype = x.dtype

    r_dtype = (
        self.router_ln.weight.dtype
        if hasattr(self.router_ln, "weight") and self.router_ln.weight is not None
        else target_dtype
    )
    normed_x = self.router_ln(x.to(r_dtype)).to(target_dtype)
    rw_dtype = (
        self.router.weight.dtype
        if hasattr(self.router, "weight") and self.router.weight is not None
        else target_dtype
    )
    logits = self.router(normed_x.to(rw_dtype)).to(target_dtype)
    if self.training and self.noise_std > 0:
        noise = torch.randn_like(logits) * self.noise_std
        scores = logits + noise
    else:
        scores = logits

    full_probs = F.softmax(scores.float(), dim=-1).to(target_dtype)
    topk_weights, topk_indices = torch.topk(full_probs, self.top_k, dim=-1)
    topk_weights = (topk_weights / topk_weights.sum(dim=-1, keepdim=True).clamp(min=1e-9)).to(
        target_dtype
    )

    router_probs = F.softmax(logits.float(), dim=-1).view(-1, self.n_experts)
    expert_mask = F.one_hot(topk_indices, num_classes=self.n_experts).float()
    expert_assignment = expert_mask.sum(dim=2).clamp(max=1.0)
    density = expert_assignment.view(-1, self.n_experts).mean(dim=0)
    prob_mean = router_probs.mean(dim=0)
    aux_loss = self.n_experts * torch.sum(density * prob_mean)
    self._aux_loss = aux_loss.detach()

    experts_stack = torch.stack([exp(x).to(target_dtype) for exp in self.experts], dim=2)
    topk_expanded = topk_indices.unsqueeze(-1).expand(-1, -1, -1, D)
    gathered = torch.gather(experts_stack, 2, topk_expanded)
    out = torch.sum(gathered * topk_weights.unsqueeze(-1), dim=2).to(target_dtype)

    return out, aux_loss

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class REAPMoEAttention (config: Any | None = None,
rank: int | None = 64,
n_experts: int = 2,
top_k: int = 1,
pretrained_attn: nn.Module | None = None,
is_decoder: bool | None = None,
pre_ln: bool | None = None,
comp_config: Any | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("reap_moe")
@ATTENTION_REGISTRY.register("reap_moe_attn")
@ATTENTION_REGISTRY.register("reap_moe_attention")
class REAPMoEAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        n_experts: int = 2,
        top_k: int = 1,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        pre_ln: bool | None = None,
        comp_config: Any | None = None,  # RP-FIX-21
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank or 64)
        self.config = config
        # If comp_config passed, prefer it; else fall back to config attribute.
        if comp_config is None:
            comp_config = getattr(config, "compression_config", None)
        self.comp_config = comp_config

        # RP-FIX-10: strict hidden_size
        h_size = (
            kwargs.get("hidden_size")
            or getattr(config, "hidden_size", None)
            or getattr(config, "d_model", None)
        )
        if h_size is None and pretrained_attn is not None:
            for cand in ("q_lin", "query", "q_proj", "c_attn"):
                mod = getattr(pretrained_attn, cand, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    if w.ndim == 2:
                        if mod.__class__.__name__ == "Conv1D" or hasattr(mod, "nf"):
                            h_size = w.shape[0]
                        else:
                            h_size = w.shape[1]
                        break
        if h_size is None:
            raise ValueError("REAPMoEAttention: could not infer hidden_size.")

        self.dim = int(h_size)
        self.num_heads = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or getattr(config, "num_attention_heads", None)
            or getattr(config, "num_heads", None)
            or 12
        )
        self.num_kv_heads = (
            kwargs.get("num_key_value_heads")
            or kwargs.get("num_kv_heads")
            or getattr(config, "num_key_value_heads", None)
            or getattr(config, "num_kv_heads", None)
            or self.num_heads
        )
        # RP-FIX-23
        self.num_kv_heads = min(self.num_kv_heads, self.num_heads)
        self.head_dim = (
            kwargs.get("head_dim")
            or getattr(config, "head_dim", None)
            or getattr(config, "d_kv", None)
            or (self.dim // max(1, self.num_heads))
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        # RP-FIX-11
        self.is_decoder = _detect_decoder(config, is_decoder, pretrained_attn)
        self.is_nested_encoder = False  # RP-FIX-22

        # RP-FIX-12
        if pre_ln is None:
            pre_ln = getattr(config, "pre_ln", True) if config is not None else True
        self.pre_ln = bool(pre_ln)

        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa = True

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        target_dtype = kwargs.get("dtype")
        if target_dtype is None and source_attn is not None:
            p = next(source_attn.parameters(), None)
            if p is not None:
                target_dtype = p.dtype

        self.q_lin = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.k_lin = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.v_lin = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )
        self.out_lin = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            dtype=target_dtype,
        )

        # Aliases
        self.q_proj = self.q_lin
        self.k_proj = self.k_lin
        self.v_proj = self.v_lin
        self.out_proj = self.out_lin
        self.query = self.q_lin
        self.key = self.k_lin
        self.value = self.v_lin
        self.dense = self.out_lin
        self.o_proj = self.out_lin
        self.c_proj = self.out_lin

        self.moe = REAPMoE(
            self.dim,
            n_experts=n_experts,
            top_k=top_k,
            config=config,
            dtype=target_dtype,
        )
        self.norm1 = nn.LayerNorm(self.dim, dtype=target_dtype)
        self.norm2 = nn.LayerNorm(self.dim, dtype=target_dtype)

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        self.register_buffer("_attn_aux_loss", torch.tensor(0.0), persistent=False)

        if source_attn is not None:
            self.copy_weights(source_attn)

        if target_dtype is not None:
            self.to(target_dtype)

    @property
    def aux_loss(self) -> torch.Tensor:
        return self.moe.aux_loss

    # ------------------------------------------------------------------
    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),
        }

        # GPT-2 fused
        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()
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        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])
                    ):
                        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,
                    )

        # RoBERTa / BERT 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,
                )

        # Flat modules
        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

        # UniversalAdapter fallback
        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"REAPMoEAttention.copy_weights: missing mandatory projections "
                f"{missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_lin,
            "k": self.k_lin,
            "v": self.v_lin,
            "o": self.out_lin,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        # For nested, out_lin will be Identity (set below).
                        continue
                    else:
                        raise RuntimeError(f"REAPMoEAttention: missing weight for {proj_key}.")

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                    and w.shape[1] == target_mod.weight.shape[1]
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                norm = float(torch.linalg.norm(target_mod.weight.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "REAPMoEAttention: weight for '%s' has near-zero Frobenius norm (%.2e).",
                        proj_key,
                        norm,
                    )

        if self.freeze_base:
            for mod in [self.q_lin, self.k_lin, self.v_lin, self.out_lin]:
                mod.freeze_base_weights()

        # RP-FIX-16: nested encoder should NOT apply its own out_lin.
        if self.is_nested_encoder:
            self.out_lin = nn.Identity()

    # ------------------------------------------------------------------
    # Mask helper (RP-FIX-20)
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N_kv: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 3:
            # (B, N_q, N_kv) or (B, 1, N_kv)
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            # Reduce over query dimension if present, else keep as is.
            if v.shape[1] > 1:
                v = v.any(dim=1)  # (B, N_kv)
            else:
                v = v.squeeze(1)  # (B, N_kv)
            return v
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # 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]:
        # RP-FIX-24
        if encoder_hidden_states is not None:
            raise NotImplementedError("REAPMoEAttention does not support cross-attention.")

        B, L, _ = hidden_states.shape
        device = hidden_states.device

        target_dtype = hidden_states.dtype

        # Pre-LN: apply norm before attention (RP-FIX-17)
        if self.pre_ln:
            n1_dtype = (
                self.norm1.weight.dtype
                if hasattr(self.norm1, "weight") and self.norm1.weight is not None
                else target_dtype
            )
            h_norm = self.norm1(hidden_states.to(n1_dtype)).to(target_dtype)
            attn_input = h_norm
            residual = hidden_states
        else:
            attn_input = hidden_states
            residual = hidden_states

        # Projections on attn_input
        q = self.q_lin(attn_input).view(B, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k_new = (
            self.k_lin(attn_input).view(B, L, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
        )
        v_new = (
            self.v_lin(attn_input).view(B, L, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
        )

        # RoPE with past_len
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            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:
                    # RP-FIX-19: use layer_idx
                    idx = min(layer_idx, len(past_key_value.key_cache) - 1)
                    past_len_rope = past_key_value.key_cache[idx].shape[-2]
            cos = cos[..., past_len_rope : past_len_rope + L, :]
            sin = sin[..., past_len_rope : past_len_rope + L, :]
            q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin)

        # KV-cache
        past_k = past_v = None
        if past_key_value is not None:
            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:
            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"):
            past_key_value.update(k_new, v_new, layer_idx)

        N_kv = K_cat.shape[2]
        past_len = N_kv - L if self.is_decoder else 0

        # GQA expansion
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            K_full = repeat_kv(K_cat, repeats)
            V_full = repeat_kv(V_cat, repeats)
        else:
            K_full = K_cat
            V_full = V_cat

        # RP-FIX-18: disable SDPA if causal and past_len > 0
        use_sdpa = (
            not output_attentions and head_mask is None and not (self.is_decoder and past_len > 0)
        )

        if use_sdpa:
            q_sdp = q.to(torch.float16) if q.dtype == torch.float32 else q
            k_sdp = K_full.to(torch.float16) if K_full.dtype == torch.float32 else K_full
            v_sdp = V_full.to(torch.float16) if V_full.dtype == torch.float32 else V_full

            attn_bias_sdp = None
            if attention_mask is not None:
                if attention_mask.dim() == 4:
                    attn_bias_sdp = attention_mask.to(q_sdp.dtype)
                else:
                    token_valid = self._prepare_valid_mask(attention_mask, B, N_kv, device)
                    if token_valid is not None:
                        attn_bias_sdp = ((token_valid[:, None, None, :].float() - 1.0) * 1e4).to(
                            q_sdp.dtype
                        )

            is_causal_sdp = self.is_decoder and L > 1 and past_len == 0

            attn_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,
            )
            attn_out = (
                attn_out_sdp.to(hidden_states.dtype)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.all_head_size)
            )
            attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
            attn_ret = None
        else:
            scale = 1.0 / math.sqrt(self.head_dim)
            scores = torch.matmul(q.float(), K_full.float().transpose(-2, -1).contiguous()) * scale

            token_valid = self._prepare_valid_mask(attention_mask, B, N_kv, device)

            if self.is_decoder:
                if past_len > 0:
                    abs_q = (past_len + torch.arange(L, device=device)).unsqueeze(1)
                    abs_k = torch.arange(N_kv, device=device).unsqueeze(0)
                    causal = (abs_k <= abs_q).unsqueeze(0).unsqueeze(0)
                else:
                    causal = (
                        torch.tril(torch.ones(L, N_kv, device=device, dtype=torch.bool))
                        .unsqueeze(0)
                        .unsqueeze(0)
                    )
                scores = scores.masked_fill(~causal, -1e4)

            if token_valid is not None:
                scores = scores.masked_fill(~token_valid[:, None, None, :].bool(), -1e4)

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(hidden_states.dtype)
            probs = torch.nan_to_num(probs, nan=0.0)
            probs = self.dropout(probs)

            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)
                elif head_mask.dim() == 3:
                    head_mask = head_mask[:, :, :1].unsqueeze(-1)
                probs = probs * head_mask.to(probs.dtype)

            attn_out = (
                torch.matmul(probs, V_full)
                .permute(0, 2, 1, 3)
                .contiguous()
                .view(B, L, self.all_head_size)
            )
            attn_ret = probs if output_attentions else None

        # RP-FIX-16: skip out_lin for nested encoder
        if self.is_nested_encoder:
            proj = attn_out
        else:
            proj = self.out_lin(attn_out)

        proj = torch.nan_to_num(proj, nan=0.0, posinf=1e4, neginf=-1e4)

        # Residual + MoE (RP-FIX-17)
        n2_dtype = (
            self.norm2.weight.dtype
            if hasattr(self.norm2, "weight") and self.norm2.weight is not None
            else target_dtype
        )
        if self.pre_ln:
            h1 = (residual + self.dropout(proj.to(target_dtype))).to(target_dtype)
            h1_norm = self.norm2(h1.to(n2_dtype)).to(target_dtype)
            moe_out, _ = self.moe(h1_norm)
            out = (h1 + self.dropout(moe_out.to(target_dtype))).to(target_dtype)
        else:
            # Post-LN: norm after residual
            n1_dtype = (
                self.norm1.weight.dtype
                if hasattr(self.norm1, "weight") and self.norm1.weight is not None
                else target_dtype
            )
            h1 = (residual + self.dropout(proj.to(target_dtype))).to(target_dtype)
            h1 = self.norm1(h1.to(n1_dtype)).to(target_dtype)
            moe_out, _ = self.moe(h1)
            out = (h1 + self.dropout(moe_out.to(target_dtype))).to(target_dtype)
            out = self.norm2(out.to(n2_dtype)).to(target_dtype)

        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
        return out.to(target_dtype), attn_ret

Every 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.

Instance variables

prop aux_loss : torch.Tensor

Expand source code

@property
def aux_loss(self) -> torch.Tensor:
    return self.moe.aux_loss

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class RPCAttention (config: Any = None,
rank: int = 64,
pretrained_attn: nn.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("rpc_class")
@ATTENTION_REGISTRY.register("rpc_attention_class")
class RPCAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any = None,
        rank: int = 64,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank if rank is not None else 64)
        except TypeError:
            super().__init__(config)

        comp_config = kwargs.get("comp_config", getattr(config, "compression_config", None))

        self.dim = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.num_heads = int(
            _cfg_get(config, "num_attention_heads", kwargs.get("num_attention_heads", 12))
        )
        self.num_kv_heads = int(
            _cfg_get(
                config,
                "num_key_value_heads",
                _cfg_get(
                    config,
                    "num_kv_heads",
                    kwargs.get("num_key_value_heads", self.num_heads),
                ),
            )
        )

        cfg_head_dim = _cfg_get(config, "head_dim", None)
        kw_head_dim = kwargs.get("head_dim")
        inferred = cfg_head_dim or kw_head_dim

        if inferred is None and pretrained_attn is not None:
            for attr in ("q_proj", "query", "q_lin"):
                if hasattr(pretrained_attn, attr):
                    layer_q = getattr(pretrained_attn, attr)
                    if hasattr(layer_q, "weight"):
                        inferred = layer_q.weight.shape[0] // max(1, self.num_heads)
                        break

        self.head_dim = int(
            inferred or (self.dim // max(1, self.num_heads) if self.num_heads > 0 else 64)
        )
        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))

        # ---- sketch rank ----
        raw_rank = kwargs.pop(
            "rank",
            _cfg_get(
                comp_config,
                "rank",
                _cfg_get(config, "rank", rank if rank is not None else 64),
            ),
        )
        self.sketch_rank = max(1, min(int(raw_rank), self.head_dim))

        # ---- projection rank for FactorizedLoRALinear ----
        raw_proj_rank = int(
            kwargs.get(
                "proj_rank",
                _cfg_get(comp_config, "proj_rank", _cfg_get(config, "proj_rank", 64)),
            )
        )
        # RPC-13: clamp and log
        self.proj_rank = max(1, raw_proj_rank)

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            mt = str(_cfg_get(config, "model_type", "") or "").lower()
            self.is_decoder = (
                bool(_cfg_get(config, "is_decoder", False))
                or bool(kwargs.get("is_causal", False))
                or any(
                    mt.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                or (pretrained_attn is not None and hasattr(pretrained_attn, "c_attn"))
            )
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # ---- LoRA config ----
        self.lora_rank = int(
            kwargs.get(
                "lora_rank",
                _cfg_get(comp_config, "lora_rank", _cfg_get(config, "lora_rank", 16)),
            )
        )
        self.lora_alpha = float(
            kwargs.get(
                "lora_alpha",
                _cfg_get(comp_config, "lora_alpha", _cfg_get(config, "lora_alpha", 32.0)),
            )
        )

        # ---- projections ----
        FF = FactorizedLoRALinear
        self.W_q = FF(
            self.dim,
            self.all_head_size,
            rank=self.proj_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )
        self.W_k = 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.W_v = 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.W_q
        self.k_proj = self.W_k
        self.v_proj = self.W_v
        self.query = self.W_q
        self.key = self.W_k
        self.value = self.W_v

        # ---- RPC-14: seeded random projection buffers ----
        # RPC-19: non-persistent (recomputed on load with same seed)
        seed = int(kwargs.get("sketch_seed", 0xC0FFEE))
        gen = torch.Generator(device="cpu").manual_seed(seed)
        std = self.head_dim**-0.5
        proj_q = torch.randn(self.num_heads, self.head_dim, self.sketch_rank, generator=gen) * std
        proj_k = torch.randn(self.num_heads, self.head_dim, self.sketch_rank, generator=gen) * std
        self.register_buffer("proj_q", proj_q.float(), persistent=False)
        self.register_buffer("proj_k", proj_k.float(), persistent=False)

        # ---- RPC-15: configurable sketch scale ----
        # Default: 1/sqrt(head_dim) — matches the original attention scale
        # because sketching preserves the expected inner product.
        default_scale = 1.0 / math.sqrt(self.head_dim)
        self.sketch_scale = float(
            kwargs.get(
                "sketch_scale",
                _cfg_get(comp_config, "sketch_scale", default_scale),
            )
        )

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        # ---- RPC-18: storage accounting ----
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

        logger.info(
            "RPCAttention: hidden=%d, heads=%d, head_dim=%d, sketch_rank=%d, "
            "proj_rank=%d, is_decoder=%s, lora_rank=%d, scale=%.4f",
            self.dim,
            self.num_heads,
            self.head_dim,
            self.sketch_rank,
            self.proj_rank,
            self.is_decoder,
            self.lora_rank,
            self.sketch_scale,
        )

    # ------------------------------------------------------------------
    @property
    def dense(self):
        return self.out_proj

    @property
    def c_proj(self):
        return self.out_proj

    @property
    def o_proj(self):
        return self.out_proj

    @property
    def q_lin(self):
        return self.W_q

    @property
    def k_lin(self):
        return self.W_k

    @property
    def v_lin(self):
        return self.W_v

    @property
    def out_lin(self):
        return self.out_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("RPC._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_()

            # RPC-17: freeze base after SVD init.
            proj.freeze_base_weights()

    # ------------------------------------------------------------------
    # Weight copy
    # ------------------------------------------------------------------
    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)

        # RPC-18: track original storage
        try:
            for p in src.parameters():
                self._orig_storage_bytes += p.numel() * p.element_size()
        except Exception:
            pass

        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # GPT-2 fused c_attn
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c = src.c_attn
            base_c = getattr(c, "base_layer", c)
            c_w = getattr(base_c, "weight", getattr(c, "weight", None))
            if c_w is not None:
                c_w = c_w.data.clone()
                is_conv = (
                    c.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                )
                if is_conv:
                    c_w = c_w.t().contiguous()
                split = c_w.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                c_b = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (c_w[:split], c_b[:split] if c_b is not None else None)
                extracted["k"] = (
                    c_w[split : 2 * split],
                    c_b[split : 2 * split] if c_b is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :],
                    c_b[2 * split :] if c_b is not None else None,
                )

            if hasattr(src, "c_proj") and src.c_proj is not None:
                o = src.c_proj
                base_o = getattr(o, "base_layer", o)
                o_w = getattr(base_o, "weight", getattr(o, "weight", None))
                if o_w is not None:
                    o_w = o_w.data.clone()
                    is_conv_o = (
                        o.__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()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (o_w, bo.data.clone() if bo is not None else None)

        # BERT / RoBERTa nested
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    extracted[key] = (
                        m.weight.data.clone(),
                        m.bias.data.clone() if m.bias is not None else None,
                    )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                om = src.output.dense
                extracted["o"] = (
                    om.weight.data.clone(),
                    om.bias.data.clone() if om.bias is not None else None,
                )

        # Flat modules
        else:
            cand = {
                "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 key, names in cand.items():
                for n in names:
                    if hasattr(src, n):
                        m = getattr(src, n)
                        if m is not None and hasattr(m, "weight"):
                            extracted[key] = (
                                m.weight.data.clone(),
                                m.bias.data.clone() if m.bias is not None else None,
                            )
                            break

        # UniversalAdapter fallback
        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

                ua = UniversalAdapter()
                for k in missing:
                    w, b = ua.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 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"RPCAttention.copy_weights: missing {missing_after} from {cls_name}."
            )

        # ---- SVD init for each projection ----
        for key, target in (
            ("q", self.W_q),
            ("k", self.W_k),
            ("v", self.W_v),
            ("o", self.out_proj),
        ):
            w, b = extracted[key]
            if w is None:
                if key == "o" and self.is_nested_encoder:
                    eye = torch.eye(
                        target.out_features,
                        target.in_features,
                        device=target.down.weight.device,
                        dtype=target.down.weight.dtype,
                    )
                    self._svd_init(eye, None, target)
                continue

            # GQA handling for k, v
            if key in ("k", "v"):
                out_t = target.out_features
                src_out = w.shape[0]
                if src_out > out_t and src_out % out_t == 0:
                    rep = src_out // out_t
                    w = w.view(out_t, rep, -1).mean(dim=1)
                    if b is not None:
                        b = b.view(out_t, rep).mean(dim=1)
                elif src_out < out_t and out_t % src_out == 0:
                    rep = out_t // src_out
                    w = w.repeat(rep, 1)
                    if b is not None:
                        b = b.repeat(rep)

            self._svd_init(w, b, target)

        # ---- storage accounting ----
        try:
            for m in (self.W_q, self.W_k, self.W_v, self.out_proj):
                self._comp_storage_bytes += m.storage_bytes()
            for buf in (self.proj_q, self.proj_k):
                self._comp_storage_bytes += buf.numel() * buf.element_size()
        except Exception:
            pass

        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)

    # ------------------------------------------------------------------
    # RPC-06 / RPC-20: robust mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        L: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        """
        Returns a bool mask of shape (B, L) where True = valid token.
        If the input mask cannot be interpreted, returns None (no masking).
        """
        if attention_mask is None:
            return None

        def _to_bool(t: torch.Tensor) -> torch.Tensor:
            if t.dtype == torch.bool:
                return t
            if t.dtype.is_floating_point:
                return t > -1e4
            return t > 0

        if attention_mask.dim() == 2:
            m = _to_bool(attention_mask)
            if m.shape[-1] == L:
                return m
            return None
        if attention_mask.dim() == 3:
            m = _to_bool(attention_mask).squeeze(1)
            if m.shape[-1] == L:
                return m
            return None
        if attention_mask.dim() == 4:
            m = _to_bool(attention_mask).any(dim=-2).squeeze(1)
            if m.shape[-1] == L:
                return m
            return None

        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # 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,
    ):
        B, L, D = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        L_kv = kv_input.shape[1]

        # ---- Projections ----
        q = self.W_q(hidden_states).view(B, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k = self.W_k(kv_input).view(B, L_kv, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
        v = self.W_v(kv_input).view(B, L_kv, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)

        # ---- RPC-08: RoPE with past_len ----
        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 = 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]
            cos = cos[..., past_len : past_len + L, :]
            sin = sin[..., past_len : past_len + L, :]
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # ---- GQA ----
        if self.kv_repeat > 1:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        # ---- Sketching ----
        proj_q = self.proj_q.to(dtype=q.dtype, device=q.device)
        proj_k = self.proj_k.to(dtype=k.dtype, device=k.device)

        # (B, H, L, head_dim) @ (H, head_dim, r) -> (B, H, L, r)
        q_sketched = torch.einsum("bhld,hdr->bhlr", q, proj_q)
        k_sketched = torch.einsum("bhld,hdr->bhlr", k, proj_k)

        scores = (
            torch.matmul(
                q_sketched.float(),
                k_sketched.float().transpose(-2, -1).contiguous(),
            )
            * self.sketch_scale
        )

        # ---- RPC-07 / RPC-21: causal only for self-attention ----
        if self.is_decoder and not is_cross:
            diag = 1 + max(0, L_kv - L)
            causal = (
                torch.triu(
                    torch.full((L, L_kv), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=diag,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        # ---- Attention mask ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        valid_mask = self._prepare_valid_mask(eff_mask, B, L_kv, device)
        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(dtype)
        probs = torch.nan_to_num(probs, nan=0.0)
        probs = self.dropout(probs)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)
            probs = probs * head_mask.to(probs.dtype)

        context = torch.matmul(probs, v)
        context = context.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # ---- Output projection (RPC-10: assert shape) ----
        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_proj(context)

        if out.shape[-1] != D:
            raise RuntimeError(f"RPCAttention: output feature dim {out.shape[-1]} != input {D}")

        return out, (probs if output_attentions else None)

    # ------------------------------------------------------------------
    def storage_bytes(self) -> dict[str, int]:
        """Report per-component storage."""

        def _sz(p):
            if p is None:
                return 0
            if isinstance(p, nn.Parameter):
                return p.numel() * p.element_size()
            if isinstance(p, torch.Tensor):
                return p.numel() * p.element_size()
            return 0

        out = {}
        for name, m in (
            ("W_q", self.W_q),
            ("W_k", self.W_k),
            ("W_v", self.W_v),
            ("out_proj", self.out_proj),
        ):
            out[name] = m.storage_bytes()
        out["proj_q"] = self.proj_q.numel() * self.proj_q.element_size()
        out["proj_k"] = self.proj_k.numel() * self.proj_k.element_size()
        out["total"] = sum(out.values())
        return out

Every 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.

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.W_k

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.W_q

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.W_v

Methods

def storage_bytes(self) ‑> dict[str, int]

Expand source code

def storage_bytes(self) -> dict[str, int]:
    """Report per-component storage."""

    def _sz(p):
        if p is None:
            return 0
        if isinstance(p, nn.Parameter):
            return p.numel() * p.element_size()
        if isinstance(p, torch.Tensor):
            return p.numel() * p.element_size()
        return 0

    out = {}
    for name, m in (
        ("W_q", self.W_q),
        ("W_k", self.W_k),
        ("W_v", self.W_v),
        ("out_proj", self.out_proj),
    ):
        out[name] = m.storage_bytes()
    out["proj_q"] = self.proj_q.numel() * self.proj_q.element_size()
    out["proj_k"] = self.proj_k.numel() * self.proj_k.element_size()
    out["total"] = sum(out.values())
    return out

Report per-component storage.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class ReGLAAttention (config: Any | None = None,
rank: int = 64,
kernel_size: int = 3,
pretrained_attn: torch.nn.modules.module.Module | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
is_decoder: bool | None = None,
chunk_size: int = 64,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("regla_class")
@ATTENTION_REGISTRY.register("regla_attention_class")
class ReGLAAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 64,
        kernel_size: int = 3,
        pretrained_attn: nn.Module | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        is_decoder: bool | None = None,
        chunk_size: int = 64,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        comp_config = getattr(config, "compression_config", None) if config else None
        src = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")

        # -------- hidden_size --------
        dim = hidden_size
        if dim is None:
            dim = _cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768))
        self.dim = int(dim)

        # -------- num_heads --------
        n_heads = num_heads
        if n_heads is None:
            n_heads = _cfg_get(config, "num_attention_heads", kwargs.get("num_attention_heads", 12))
        if n_heads is None and src is not None:
            n_heads = getattr(
                src,
                "num_heads",
                getattr(src, "num_attention_heads", getattr(src, "n_head", None)),
            )
        if n_heads is None:
            n_heads = max(1, self.dim // 64)
        self.num_heads = max(1, int(n_heads))

        # -------- head_dim --------
        cfg_hd = _cfg_get(config, "head_dim", None)
        if cfg_hd is not None:
            self.head_dim = int(cfg_hd)
        elif src is not None:
            self.head_dim = None
            for attr in ("q_proj", "query", "q_lin"):
                if hasattr(src, attr):
                    qw = getattr(src, attr)
                    w = getattr(qw, "weight", None)
                    if w is not None:
                        self.head_dim = int(w.shape[0] // self.num_heads)
                        break
            if self.head_dim is None:
                self.head_dim = self.dim // self.num_heads
        else:
            self.head_dim = self.dim // self.num_heads

        if self.dim % self.num_heads != 0:
            self.head_dim = self.dim // self.num_heads

        # -------- num_kv_heads --------
        self.num_kv_heads = int(
            _cfg_get(
                config,
                "num_key_value_heads",
                kwargs.get("num_key_value_heads", self.num_heads),
            )
        )
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = max(1, self.num_heads // self.num_kv_heads)

        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        self.kernel_size = int(kernel_size)
        self.chunk_size = max(1, int(chunk_size))

        self.conv_scale = nn.Parameter(torch.zeros(1))

        # -------- is_decoder --------
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            mt = str(_cfg_get(config, "model_type", "") or "").lower()
            self.is_decoder = (
                bool(_cfg_get(config, "is_decoder", False))
                or bool(_cfg_get(config, "causal", False))
                or any(
                    mt.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                    )
                )
                or (src is not None and hasattr(src, "c_attn"))
                or bool(kwargs.get("is_causal", False))
            )
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        # -------- LoRA config --------
        self.lora_rank = int(
            _cfg_get(
                config,
                "lora_rank",
                _cfg_get(comp_config, "lora_rank", kwargs.get("lora_rank", 8)),
            )
        )
        self.lora_alpha = float(
            _cfg_get(
                config,
                "lora_alpha",
                _cfg_get(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        )
        self.freeze_base = self.lora_rank > 0

        # -------- device/dtype --------
        target_dtype = kwargs.get("dtype")
        target_device = kwargs.get("device")
        if target_dtype is None and src is not None:
            p = next(src.parameters(), None)
            if p is not None:
                target_dtype = p.dtype
                if target_device is None:
                    target_device = p.device
        if target_dtype is None:
            target_dtype = torch.float32

        # -------- projections --------
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            device=target_device,
            dtype=target_dtype,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            device=target_device,
            dtype=target_dtype,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            device=target_device,
            dtype=target_dtype,
        )
        self.gate_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            device=target_device,
            dtype=target_dtype,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
            device=target_device,
            dtype=target_dtype,
        )

        # Aliases
        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj
        self.q_lin = self.q_proj
        self.k_lin = self.k_proj
        self.v_lin = self.v_proj
        self.out_lin = self.out_proj
        self.dense = self.out_proj

        # Conv1d
        self.conv = nn.Conv1d(
            in_channels=self.dim,
            out_channels=self.dim,
            kernel_size=self.kernel_size,
            groups=self.dim,
            bias=True,
            device=target_device,
            dtype=target_dtype,
        )

        with torch.no_grad():
            nn.init.normal_(self.gate_proj.weight, mean=0.0, std=0.02)
            if self.gate_proj.bias is not None:
                self.gate_proj.bias.zero_()

        if src is not None:
            self.copy_weights(src)

    # ------------------------------------------------------------------
    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),
        }

        def _unwrap_linear(layer: Any) -> Any:
            return getattr(layer, "base_layer", layer)

        # GPT-2 fused
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c_attn_mod = _unwrap_linear(src.c_attn)
            c_w = c_attn_mod.weight.data.clone() if hasattr(c_attn_mod, "weight") else None
            if c_w is not None:
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or hasattr(c_attn_mod, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    c_w = c_w.t().contiguous()
                split_size = c_w.size(0) // 3
                bias_mod = 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 = _unwrap_linear(src.c_proj)
                o_w = c_proj_mod.weight.data.clone() if hasattr(c_proj_mod, "weight") else None
                if o_w is not None:
                    if (
                        c_proj_mod.__class__.__name__ == "Conv1D"
                        or hasattr(c_proj_mod, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        o_w = o_w.t().contiguous()
                    bias_o = getattr(c_proj_mod, "bias", None)
                    extracted["o"] = (
                        o_w,
                        bias_o.data.clone() if bias_o is not None else None,
                    )

        # RoBERTa nested (outer container)
        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,
                )

        # Flat modules
        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.data.clone(),
                            b.data.clone() if b is not None else None,
                        )
            except Exception as e:
                logger.debug("UniversalAdapter 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"ReGLAAttention.copy_weights: missing {missing_after} from {cls_name}."
            )

        with torch.no_grad():
            for proj_key, target_mod in (
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ):
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            target_mod.bias.zero_()
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                if (
                    w.shape != target_mod.weight.shape
                    and target_mod.weight.shape == w.t().contiguous().shape
                ):
                    w = w.t().contiguous()

                out_d = min(target_mod.weight.size(0), w.size(0))
                in_d = min(target_mod.weight.size(1), w.size(1))
                target_mod.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(
                        device=target_mod.weight.device, dtype=target_mod.weight.dtype
                    )
                )

                if b is not None and target_mod.bias is not None:
                    min_b = min(target_mod.bias.numel(), b.numel())
                    target_mod.bias.data[:min_b].copy_(
                        b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype)
                    )
                elif target_mod.bias is not None:
                    target_mod.bias.data.zero_()

                w_norm = float(torch.norm(target_mod.weight.data).item())
                if w_norm < 1e-4:
                    logger.warning(
                        "ReGLAAttention: weight for '%s' has near-zero norm (%.2e)",
                        proj_key,
                        w_norm,
                    )

                if self.freeze_base:
                    target_mod.freeze_base_weights()

    # ------------------------------------------------------------------
    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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        layer_idx: int = 0,
        *args: Any,
        **kwargs: Any,
    ):
        if encoder_hidden_states is not None:
            raise NotImplementedError("ReGLAAttention does not support cross-attention")

        x = hidden_states
        B, L, D = x.shape
        device = x.device
        dtype = x.dtype

        q = (
            self.q_proj(x)
            .view(B, L, self.num_heads, self.head_dim)
            .permute(0, 2, 1, 3)
            .contiguous()
        )
        k = (
            self.k_proj(x)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
            .contiguous()
        )
        v = (
            self.v_proj(x)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
            .contiguous()
        )
        g = (
            self.gate_proj(x)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
            .contiguous()
        )

        # RoPE with past_len offset
        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]
            cos = cos[..., past_len : past_len + L, :]
            sin = sin[..., past_len : past_len + L, :]
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # GQA expansion
        if self.num_heads != self.num_kv_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)
            g = repeat_kv(g, self.kv_repeat)

        # Non-negative kernel activation & gating
        q_act = F.elu(q.float()) + 1.0
        k_act = F.elu(k.float()) + 1.0
        g_sig = torch.sigmoid(g.float())

        # Valid mask
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                valid_mask = (
                    (attention_mask > 0)
                    if not attention_mask.dtype.is_floating_point
                    else (attention_mask > -1e4)
                )
            elif attention_mask.dim() == 4:
                if attention_mask.dtype.is_floating_point:
                    valid_mask = (attention_mask > -1e4).any(dim=-2).squeeze(1)
                else:
                    valid_mask = attention_mask.bool().any(dim=-2).squeeze(1)
            else:
                valid_mask = torch.ones(B, L, device=device, dtype=torch.bool)
        else:
            valid_mask = torch.ones(B, L, device=device, dtype=torch.bool)

        mask_4d = valid_mask[:, None, :, None].to(q_act.dtype)
        k_gated = k_act * g_sig * mask_4d
        v_gated = v.float() * g_sig * mask_4d

        is_causal = self.is_decoder or kwargs.get("is_causal", False)

        if is_causal:
            chunk = self.chunk_size
            S = torch.zeros(
                B,
                self.num_heads,
                self.head_dim,
                self.head_dim,
                device=device,
                dtype=q_act.dtype,
            )
            Z = torch.zeros(B, self.num_heads, self.head_dim, device=device, dtype=q_act.dtype)
            out_list = []

            for start in range(0, L, chunk):
                end = min(start + chunk, L)
                q_c = q_act[:, :, start:end]
                k_c = k_gated[:, :, start:end]
                v_c = v_gated[:, :, start:end]

                intra_kv = torch.cumsum(k_c.unsqueeze(-1) * v_c.unsqueeze(-2), dim=2)
                num_intra = (q_c.unsqueeze(-2) @ intra_kv).squeeze(-2)
                num_inter = (q_c.unsqueeze(-2) @ S.unsqueeze(2)).squeeze(-2)
                num = num_intra + num_inter

                z_intra = torch.cumsum(k_c, dim=2)
                den = (q_c * (z_intra + Z.unsqueeze(2))).sum(-1, keepdim=True).clamp(min=1e-4)
                out_list.append(num / den)

                S = S + (k_c.unsqueeze(-1) * v_c.unsqueeze(-2)).sum(dim=2)
                Z = Z + k_c.sum(dim=2)

            out = torch.cat(out_list, dim=2)
        else:
            kv = torch.matmul(k_gated.transpose(-2, -1).contiguous(), v_gated)
            z = k_gated.sum(dim=2, keepdim=True)
            num = torch.matmul(q_act, kv)
            den = (q_act * z).sum(-1, keepdim=True).clamp(min=1e-4)
            out = num / den

        out = out * mask_4d

        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
            out = out * hm.to(out.dtype)

        out = out.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size).to(dtype)

        # Depthwise 1D conv branch
        x_in = (x * valid_mask.unsqueeze(-1).to(dtype)).permute(0, 2, 1)
        conv_w = self.conv.weight.to(dtype)
        conv_b = self.conv.bias.to(dtype) if self.conv.bias is not None else None
        if is_causal:
            x_pad = F.pad(x_in, (self.kernel_size - 1, 0))
            loc = F.conv1d(
                x_pad,
                conv_w,
                conv_b,
                stride=self.conv.stride,
                padding=self.conv.padding,
                dilation=self.conv.dilation,
                groups=self.conv.groups,
            ).permute(0, 2, 1)
        else:
            p = self.kernel_size // 2
            x_pad = F.pad(x_in, (p, p))
            loc = F.conv1d(
                x_pad,
                conv_w,
                conv_b,
                stride=self.conv.stride,
                padding=self.conv.padding,
                dilation=self.conv.dilation,
                groups=self.conv.groups,
            ).permute(0, 2, 1)

        if loc.shape[-1] != self.dim:
            raise RuntimeError(
                f"ReGLAAttention: conv output dim {loc.shape[-1]} != hidden dim {self.dim}"
            )

        fused = (out.to(dtype) + self.conv_scale.to(dtype) * loc.to(dtype)).to(dtype)
        fused = fused * valid_mask.unsqueeze(-1).to(dtype)

        if self.is_nested_encoder:
            output = fused
        else:
            output = self.out_proj(fused)

        if output.shape[-1] != D:
            raise RuntimeError(f"ReGLAAttention: output dim {output.shape[-1]} != input {D}")

        return output.to(dtype), None

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class Replace (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
num_attention_heads: int | None = None,
dropout: float = 0.1,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("replace_class")
@ATTENTION_REGISTRY.register("replace_attention_class")
class ReplaceAttention(BaseAttention):
    """
    Robust identity / low-rank SVD replacement attention.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    # RE-28: class-level flag (not per-instance) so deepcopy doesn't reset it.
    _first_copy_done: bool = False

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        num_attention_heads: int | None = None,
        dropout: float = 0.1,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank, **kwargs)
        except TypeError:
            super().__init__(config)

        self.hidden_size = int(
            hidden_size or _cfg_get(config, "hidden_size", _cfg_get(config, "d_model", 768))
        )
        n_heads_candidate = (
            num_heads
            or num_attention_heads
            or _cfg_get(config, "num_attention_heads", _cfg_get(config, "num_heads", 12))
        )
        self.num_heads = int(n_heads_candidate) if n_heads_candidate else 12
        if self.num_heads <= 0:
            self.num_heads = 12
        self.head_dim = max(1, self.hidden_size // self.num_heads)
        self.all_head_size = self.num_heads * self.head_dim

        self.num_kv_heads = int(
            _cfg_get(
                config,
                "num_key_value_heads",
                _cfg_get(config, "num_kv_heads", self.num_heads),
            )
        )
        if self.num_kv_heads <= 0 or self.num_kv_heads > self.num_heads:
            self.num_kv_heads = self.num_heads

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = bool(_cfg_get(config, "is_decoder", False))

        self.is_nested_encoder: bool = False

        self.rank = int(rank) if rank and rank > 0 else 0

        if self.rank > 0:
            self.q_proj = _FactorizedLinear(
                self.hidden_size, self.all_head_size, self.rank, bias=True
            )
            self.k_proj = _FactorizedLinear(
                self.hidden_size,
                self.num_kv_heads * self.head_dim,
                self.rank,
                bias=True,
            )
            self.v_proj = _FactorizedLinear(
                self.hidden_size,
                self.num_kv_heads * self.head_dim,
                self.rank,
                bias=True,
            )
            self.out_proj = _FactorizedLinear(
                self.all_head_size, self.hidden_size, self.rank, bias=True
            )
        else:
            self.q_proj = nn.Linear(self.hidden_size, self.all_head_size)
            self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim)
            self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim)
            self.out_proj = nn.Linear(self.all_head_size, self.hidden_size)

        self.dropout = nn.Dropout(dropout)

        if old_attn is not None:
            self.copy_weights(old_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 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 dense(self):
        return self.out_proj

    @property
    def c_proj(self):
        return self.out_proj

    # ------------------------------------------------------------------
    # RE-26: _copy_into for _FactorizedLinear
    # ------------------------------------------------------------------
    def _copy_into(
        self,
        target_proj: nn.Module,
        w: torch.Tensor | None,
        b: torch.Tensor | None,
        target_out: int,
        target_in: int,
    ) -> None:
        if w is None:
            return
        with torch.no_grad():
            if isinstance(target_proj, _FactorizedLinear):
                try:
                    W = w.detach().float()
                    if W.shape[0] != target_out or W.shape[1] != target_in:
                        logger.warning(
                            "Shape mismatch in copy_into: source %s, target (%d, %d).",
                            tuple(W.shape),
                            target_out,
                            target_in,
                        )
                    # Use up to target_proj.rank.
                    req_rank = target_proj.rank
                    if W.shape[0] > 512 and W.shape[1] > 512:
                        U, S, Vh = _randomized_svd(W, req_rank)
                    else:
                        U, S, Vh = torch.linalg.svd(W, full_matrices=False)
                    effective_rank = min(req_rank, U.size(1), Vh.size(0))
                    # RE-26: zero first, then fill only effective_rank rows/cols.
                    target_proj.down.weight.data.zero_()
                    target_proj.up.weight.data.zero_()
                    target_proj.down.weight.data[:effective_rank, :].copy_(
                        Vh[:effective_rank, :].to(target_proj.down.weight.dtype)
                    )
                    US = U[:, :effective_rank] @ torch.diag(S[:effective_rank])
                    target_proj.up.weight.data[:, :effective_rank].copy_(
                        US.to(target_proj.up.weight.dtype)
                    )
                    if b is not None and target_proj.up.bias is not None:
                        n = min(b.numel(), target_proj.up.bias.numel())
                        target_proj.up.bias.data[:n].copy_(b[:n].to(target_proj.up.bias.dtype))
                except Exception as e:
                    logger.error("SVD init failed: %s", e)
            else:
                if w.shape != target_proj.weight.shape:
                    raise ValueError(
                        f"Weight shape mismatch: source {tuple(w.shape)} != "
                        f"target {tuple(target_proj.weight.shape)}."
                    )
                target_proj.weight.data.copy_(w.to(target_proj.weight.dtype))
                if b is not None and target_proj.bias is not None:
                    if b.shape != target_proj.bias.shape:
                        raise ValueError(
                            f"Bias shape mismatch: {tuple(b.shape)} != "
                            f"{tuple(target_proj.bias.shape)}"
                        )
                    target_proj.bias.data.copy_(b.to(target_proj.bias.dtype))

    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        if old_attn is None:
            return

        src = getattr(old_attn, "base_layer", old_attn)
        w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

        # Pattern 1: GPT-2 Conv1D / Linear fused
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn = src.c_attn
            c_proj = src.c_proj
            w_all = c_attn.weight.data
            b_all = c_attn.bias.data if c_attn.bias is not None else None

            hidden = w_all.shape[0] if w_all.shape[0] < w_all.shape[1] else w_all.shape[1]
            if w_all.shape[0] == 3 * hidden or w_all.shape[1] == 3 * hidden:
                if w_all.shape[1] == 3 * hidden:
                    chunks = w_all.chunk(3, dim=1)
                    w_q, w_k, w_v = [c.t().contiguous() for c in chunks]
                else:
                    chunks = w_all.chunk(3, dim=0)
                    w_q, w_k, w_v = chunks[0], chunks[1], chunks[2]
                if b_all is not None:
                    b_q, b_k, b_v = b_all.chunk(3, dim=0)

            w_proj = c_proj.weight.data
            if w_proj.shape[0] != self.hidden_size and w_proj.shape[1] == self.hidden_size:
                w_o = w_proj.t().contiguous()
            else:
                w_o = w_proj
            b_o = c_proj.bias.data if c_proj.bias is not None else None

        # Pattern 2: BERT / RoBERTa nested
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            w_q = s.query.weight.data
            b_q = s.query.bias.data if s.query.bias is not None else None
            w_k = s.key.weight.data
            b_k = s.key.bias.data if s.key.bias is not None else None
            w_v = s.value.weight.data
            b_v = s.value.bias.data if s.value.bias is not None else None
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                w_o = src.output.dense.weight.data
                b_o = src.output.dense.bias.data if src.output.dense.bias is not None else None
            self.is_nested_encoder = True
            self.out_proj = nn.Identity()

        # Pattern 3: Flat query/key/value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.data
            b_q = src.query.bias.data if src.query.bias is not None else None
            w_k = src.key.weight.data
            b_k = src.key.bias.data if src.key.bias is not None else None
            w_v = src.value.weight.data
            b_v = src.value.bias.data if src.value.bias is not None else None
            for attr in ("dense", "out_proj", "o_proj", "proj"):
                if hasattr(src, attr):
                    proj = getattr(src, attr)
                    if hasattr(proj, "weight"):
                        w_o = proj.weight.data
                        b_o = proj.bias.data if getattr(proj, "bias", None) is not None else None
                    break
            if w_o is None:
                self.is_nested_encoder = True
                self.out_proj = nn.Identity()

        # Pattern 4: DistilBERT
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.data
            b_q = src.q_lin.bias.data if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.data
            b_k = src.k_lin.bias.data if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.data
            b_v = src.v_lin.bias.data if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.data
                b_o = src.out_lin.bias.data if src.out_lin.bias is not None else None

        # Pattern 5: Qwen / LLaMA
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.data
            b_q = src.q_proj.bias.data if getattr(src.q_proj, "bias", None) is not None else None
            w_k = src.k_proj.weight.data
            b_k = src.k_proj.bias.data if getattr(src.k_proj, "bias", None) is not None else None
            w_v = src.v_proj.weight.data
            b_v = src.v_proj.bias.data if getattr(src.v_proj, "bias", None) is not None else None
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.data
                b_o = (
                    src.o_proj.bias.data if getattr(src.o_proj, "bias", None) is not None else None
                )

        # Fallback
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                if w_q is None:
                    w_q, b_q = adapter.get_weight_bias(old_attn, "q")
                if w_k is None:
                    w_k, b_k = adapter.get_weight_bias(old_attn, "k")
                if w_v is None:
                    w_v, b_v = adapter.get_weight_bias(old_attn, "v")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(old_attn, "o")
            except Exception as e:
                logger.warning("UniversalAdapter fallback failed: %s", e)

        # RE-28: use class-level flag
        if not ReplaceAttention._first_copy_done:
            for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v), ("w_o", w_o)]:
                if w is not None:
                    norm = float(torch.linalg.norm(w.float()).item())
                    if norm < 1e-4:
                        logger.warning(
                            "Copied weight %s has near-zero Frobenius norm (%.2e)",
                            name,
                            norm,
                        )
            ReplaceAttention._first_copy_done = True

        # RE-27: cast only the affected projections, not the whole module.
        if w_q is not None and hasattr(self.q_proj, "weight"):
            if w_q.dtype != self.q_proj.weight.dtype:
                self.q_proj.to(w_q.dtype)
        if w_k is not None and hasattr(self.k_proj, "weight"):
            if w_k.dtype != self.k_proj.weight.dtype:
                self.k_proj.to(w_k.dtype)
        if w_v is not None and hasattr(self.v_proj, "weight"):
            if w_v.dtype != self.v_proj.weight.dtype:
                self.v_proj.to(w_v.dtype)
        if w_o is not None and not self.is_nested_encoder and hasattr(self.out_proj, "weight"):
            if w_o.dtype != self.out_proj.weight.dtype:
                self.out_proj.to(w_o.dtype)

        self._copy_into(self.q_proj, w_q, b_q, self.all_head_size, self.hidden_size)
        self._copy_into(self.k_proj, w_k, b_k, self.num_kv_heads * self.head_dim, self.hidden_size)
        self._copy_into(self.v_proj, w_v, b_v, self.num_kv_heads * self.head_dim, self.hidden_size)

        if not self.is_nested_encoder:
            self._copy_into(self.out_proj, w_o, b_o, self.hidden_size, self.all_head_size)

    # ------------------------------------------------------------------
    # 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,
        output_attentions: bool = False,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("ReplaceAttention: cross-attention not supported.")

        B, N, _ = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        if self.num_kv_heads < self.num_heads:
            repeat = self.num_heads // self.num_kv_heads
            K = K.repeat_interleave(repeat, dim=1)
            V = V.repeat_interleave(repeat, dim=1)

        is_causal = self.is_decoder and N > 1

        attn_mask_sdpa = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                attn_mask_sdpa = (
                    (attention_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
                ).to(Q.dtype)
            elif attention_mask.dim() == 4:
                attn_mask_sdpa = attention_mask.to(Q.dtype)

        if not output_attentions and head_mask is None:
            try:
                q_sdpa = Q.half() if Q.dtype == torch.float32 else Q
                k_sdpa = K.half() if K.dtype == torch.float32 else K
                v_sdpa = V.half() if V.dtype == torch.float32 else V

                out_sdpa = F.scaled_dot_product_attention(
                    q_sdpa,
                    k_sdpa,
                    v_sdpa,
                    attn_mask=attn_mask_sdpa,
                    is_causal=is_causal and attn_mask_sdpa is None,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                out = out_sdpa.to(dtype).transpose(1, 2).contiguous().view(B, N, self.all_head_size)
                out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

                if self.is_nested_encoder:
                    return out, None
                result = self.out_proj(out)
                result = torch.nan_to_num(result, nan=0.0, posinf=1e4, neginf=-1e4)
                return result, None
            except Exception as e:
                logger.debug("SDPA path failed (%s); falling back.", e)

        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if is_causal:
            causal = (
                torch.triu(
                    torch.full((N, N), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        if attention_mask is not None:
            if attention_mask.dim() == 2:
                scores = scores + (attention_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif attention_mask.dim() == 3:
                scores = scores + attention_mask.unsqueeze(1).float()
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(dtype)
        probs = self.dropout(probs)

        # RE-25: proper head_mask handling.
        if head_mask is not None:
            if head_mask.dim() == 1:
                hm = head_mask.view(1, -1, 1, 1)
            elif head_mask.dim() == 2:
                # (num_layers, num_heads) or (B, H): pick the right slice.
                if head_mask.shape[0] > 1 and head_mask.shape[1] == self.num_heads:
                    hm = head_mask[:1, :].view(1, -1, 1, 1)
                else:
                    hm = head_mask.view(head_mask.shape[0], -1, 1, 1)
            elif head_mask.dim() == 3:
                # Treat as (B, H, N) — mask per head; average over N.
                # Simpler: reduce to (B, H, 1, 1).
                hm = head_mask.mean(dim=-1).view(head_mask.shape[0], -1, 1, 1)
            elif head_mask.dim() == 4:
                hm = head_mask
            else:
                hm = None
            if hm is not None:
                probs = probs * hm.to(probs.dtype)

        context = torch.matmul(probs, V)

        valid_mask = _get_valid_mask(attention_mask, B, N, device)
        context = context * valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)

        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            return context, (probs if output_attentions else None)

        out = self.out_proj(context)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        return out, (probs if output_attentions else None)

Robust identity / low-rank SVD replacement attention.

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.

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

    src = getattr(old_attn, "base_layer", old_attn)
    w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

    # Pattern 1: GPT-2 Conv1D / Linear fused
    if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
        c_attn = src.c_attn
        c_proj = src.c_proj
        w_all = c_attn.weight.data
        b_all = c_attn.bias.data if c_attn.bias is not None else None

        hidden = w_all.shape[0] if w_all.shape[0] < w_all.shape[1] else w_all.shape[1]
        if w_all.shape[0] == 3 * hidden or w_all.shape[1] == 3 * hidden:
            if w_all.shape[1] == 3 * hidden:
                chunks = w_all.chunk(3, dim=1)
                w_q, w_k, w_v = [c.t().contiguous() for c in chunks]
            else:
                chunks = w_all.chunk(3, dim=0)
                w_q, w_k, w_v = chunks[0], chunks[1], chunks[2]
            if b_all is not None:
                b_q, b_k, b_v = b_all.chunk(3, dim=0)

        w_proj = c_proj.weight.data
        if w_proj.shape[0] != self.hidden_size and w_proj.shape[1] == self.hidden_size:
            w_o = w_proj.t().contiguous()
        else:
            w_o = w_proj
        b_o = c_proj.bias.data if c_proj.bias is not None else None

    # Pattern 2: BERT / RoBERTa nested
    elif hasattr(src, "self") and hasattr(src.self, "query"):
        s = src.self
        w_q = s.query.weight.data
        b_q = s.query.bias.data if s.query.bias is not None else None
        w_k = s.key.weight.data
        b_k = s.key.bias.data if s.key.bias is not None else None
        w_v = s.value.weight.data
        b_v = s.value.bias.data if s.value.bias is not None else None
        if hasattr(src, "output") and hasattr(src.output, "dense"):
            w_o = src.output.dense.weight.data
            b_o = src.output.dense.bias.data if src.output.dense.bias is not None else None
        self.is_nested_encoder = True
        self.out_proj = nn.Identity()

    # Pattern 3: Flat query/key/value
    elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
        w_q = src.query.weight.data
        b_q = src.query.bias.data if src.query.bias is not None else None
        w_k = src.key.weight.data
        b_k = src.key.bias.data if src.key.bias is not None else None
        w_v = src.value.weight.data
        b_v = src.value.bias.data if src.value.bias is not None else None
        for attr in ("dense", "out_proj", "o_proj", "proj"):
            if hasattr(src, attr):
                proj = getattr(src, attr)
                if hasattr(proj, "weight"):
                    w_o = proj.weight.data
                    b_o = proj.bias.data if getattr(proj, "bias", None) is not None else None
                break
        if w_o is None:
            self.is_nested_encoder = True
            self.out_proj = nn.Identity()

    # Pattern 4: DistilBERT
    elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
        w_q = src.q_lin.weight.data
        b_q = src.q_lin.bias.data if src.q_lin.bias is not None else None
        w_k = src.k_lin.weight.data
        b_k = src.k_lin.bias.data if src.k_lin.bias is not None else None
        w_v = src.v_lin.weight.data
        b_v = src.v_lin.bias.data if src.v_lin.bias is not None else None
        if hasattr(src, "out_lin"):
            w_o = src.out_lin.weight.data
            b_o = src.out_lin.bias.data if src.out_lin.bias is not None else None

    # Pattern 5: Qwen / LLaMA
    elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
        w_q = src.q_proj.weight.data
        b_q = src.q_proj.bias.data if getattr(src.q_proj, "bias", None) is not None else None
        w_k = src.k_proj.weight.data
        b_k = src.k_proj.bias.data if getattr(src.k_proj, "bias", None) is not None else None
        w_v = src.v_proj.weight.data
        b_v = src.v_proj.bias.data if getattr(src.v_proj, "bias", None) is not None else None
        if hasattr(src, "o_proj"):
            w_o = src.o_proj.weight.data
            b_o = (
                src.o_proj.bias.data if getattr(src.o_proj, "bias", None) is not None else None
            )

    # Fallback
    if w_q is None or w_k is None or w_v is None:
        try:
            adapter = UniversalAdapter()
            if w_q is None:
                w_q, b_q = adapter.get_weight_bias(old_attn, "q")
            if w_k is None:
                w_k, b_k = adapter.get_weight_bias(old_attn, "k")
            if w_v is None:
                w_v, b_v = adapter.get_weight_bias(old_attn, "v")
            if w_o is None:
                w_o, b_o = adapter.get_weight_bias(old_attn, "o")
        except Exception as e:
            logger.warning("UniversalAdapter fallback failed: %s", e)

    # RE-28: use class-level flag
    if not ReplaceAttention._first_copy_done:
        for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v), ("w_o", w_o)]:
            if w is not None:
                norm = float(torch.linalg.norm(w.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "Copied weight %s has near-zero Frobenius norm (%.2e)",
                        name,
                        norm,
                    )
        ReplaceAttention._first_copy_done = True

    # RE-27: cast only the affected projections, not the whole module.
    if w_q is not None and hasattr(self.q_proj, "weight"):
        if w_q.dtype != self.q_proj.weight.dtype:
            self.q_proj.to(w_q.dtype)
    if w_k is not None and hasattr(self.k_proj, "weight"):
        if w_k.dtype != self.k_proj.weight.dtype:
            self.k_proj.to(w_k.dtype)
    if w_v is not None and hasattr(self.v_proj, "weight"):
        if w_v.dtype != self.v_proj.weight.dtype:
            self.v_proj.to(w_v.dtype)
    if w_o is not None and not self.is_nested_encoder and hasattr(self.out_proj, "weight"):
        if w_o.dtype != self.out_proj.weight.dtype:
            self.out_proj.to(w_o.dtype)

    self._copy_into(self.q_proj, w_q, b_q, self.all_head_size, self.hidden_size)
    self._copy_into(self.k_proj, w_k, b_k, self.num_kv_heads * self.head_dim, self.hidden_size)
    self._copy_into(self.v_proj, w_v, b_v, self.num_kv_heads * self.head_dim, self.hidden_size)

    if not self.is_nested_encoder:
        self._copy_into(self.out_proj, w_o, b_o, self.hidden_size, self.all_head_size)

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,
output_attentions: bool = False,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
**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,
    output_attentions: bool = False,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    if encoder_hidden_states is not None:
        raise NotImplementedError("ReplaceAttention: cross-attention not supported.")

    B, N, _ = hidden_states.shape
    device = hidden_states.device
    dtype = hidden_states.dtype

    Q = (
        self.q_proj(hidden_states)
        .view(B, N, self.num_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    K = (
        self.k_proj(hidden_states)
        .view(B, N, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    V = (
        self.v_proj(hidden_states)
        .view(B, N, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    if self.num_kv_heads < self.num_heads:
        repeat = self.num_heads // self.num_kv_heads
        K = K.repeat_interleave(repeat, dim=1)
        V = V.repeat_interleave(repeat, dim=1)

    is_causal = self.is_decoder and N > 1

    attn_mask_sdpa = None
    if attention_mask is not None:
        if attention_mask.dim() == 2:
            attn_mask_sdpa = (
                (attention_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            ).to(Q.dtype)
        elif attention_mask.dim() == 4:
            attn_mask_sdpa = attention_mask.to(Q.dtype)

    if not output_attentions and head_mask is None:
        try:
            q_sdpa = Q.half() if Q.dtype == torch.float32 else Q
            k_sdpa = K.half() if K.dtype == torch.float32 else K
            v_sdpa = V.half() if V.dtype == torch.float32 else V

            out_sdpa = F.scaled_dot_product_attention(
                q_sdpa,
                k_sdpa,
                v_sdpa,
                attn_mask=attn_mask_sdpa,
                is_causal=is_causal and attn_mask_sdpa is None,
                dropout_p=self.dropout.p if self.training else 0.0,
            )
            out = out_sdpa.to(dtype).transpose(1, 2).contiguous().view(B, N, self.all_head_size)
            out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

            if self.is_nested_encoder:
                return out, None
            result = self.out_proj(out)
            result = torch.nan_to_num(result, nan=0.0, posinf=1e4, neginf=-1e4)
            return result, None
        except Exception as e:
            logger.debug("SDPA path failed (%s); falling back.", e)

    scale = 1.0 / math.sqrt(self.head_dim)
    scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

    if is_causal:
        causal = (
            torch.triu(
                torch.full((N, N), float("-inf"), device=device, dtype=torch.float32),
                diagonal=1,
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        scores = scores + causal

    if attention_mask is not None:
        if attention_mask.dim() == 2:
            scores = scores + (attention_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
        elif attention_mask.dim() == 3:
            scores = scores + attention_mask.unsqueeze(1).float()
        elif attention_mask.dim() == 4:
            scores = scores + attention_mask.float()

    scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
    probs = F.softmax(scores, dim=-1).to(dtype)
    probs = self.dropout(probs)

    # RE-25: proper head_mask handling.
    if head_mask is not None:
        if head_mask.dim() == 1:
            hm = head_mask.view(1, -1, 1, 1)
        elif head_mask.dim() == 2:
            # (num_layers, num_heads) or (B, H): pick the right slice.
            if head_mask.shape[0] > 1 and head_mask.shape[1] == self.num_heads:
                hm = head_mask[:1, :].view(1, -1, 1, 1)
            else:
                hm = head_mask.view(head_mask.shape[0], -1, 1, 1)
        elif head_mask.dim() == 3:
            # Treat as (B, H, N) — mask per head; average over N.
            # Simpler: reduce to (B, H, 1, 1).
            hm = head_mask.mean(dim=-1).view(head_mask.shape[0], -1, 1, 1)
        elif head_mask.dim() == 4:
            hm = head_mask
        else:
            hm = None
        if hm is not None:
            probs = probs * hm.to(probs.dtype)

    context = torch.matmul(probs, V)

    valid_mask = _get_valid_mask(attention_mask, B, N, device)
    context = context * valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)

    context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

    if self.is_nested_encoder:
        return context, (probs if output_attentions else None)

    out = self.out_proj(context)
    out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

    return out, (probs if output_attentions else None)

Must return (output, attention_weights_or_None).

class ReplaceAttention (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
num_attention_heads: int | None = None,
dropout: float = 0.1,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("replace_class")
@ATTENTION_REGISTRY.register("replace_attention_class")
class ReplaceAttention(BaseAttention):
    """
    Robust identity / low-rank SVD replacement attention.
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    # RE-28: class-level flag (not per-instance) so deepcopy doesn't reset it.
    _first_copy_done: bool = False

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        num_attention_heads: int | None = None,
        dropout: float = 0.1,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank, **kwargs)
        except TypeError:
            super().__init__(config)

        self.hidden_size = int(
            hidden_size or _cfg_get(config, "hidden_size", _cfg_get(config, "d_model", 768))
        )
        n_heads_candidate = (
            num_heads
            or num_attention_heads
            or _cfg_get(config, "num_attention_heads", _cfg_get(config, "num_heads", 12))
        )
        self.num_heads = int(n_heads_candidate) if n_heads_candidate else 12
        if self.num_heads <= 0:
            self.num_heads = 12
        self.head_dim = max(1, self.hidden_size // self.num_heads)
        self.all_head_size = self.num_heads * self.head_dim

        self.num_kv_heads = int(
            _cfg_get(
                config,
                "num_key_value_heads",
                _cfg_get(config, "num_kv_heads", self.num_heads),
            )
        )
        if self.num_kv_heads <= 0 or self.num_kv_heads > self.num_heads:
            self.num_kv_heads = self.num_heads

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = bool(_cfg_get(config, "is_decoder", False))

        self.is_nested_encoder: bool = False

        self.rank = int(rank) if rank and rank > 0 else 0

        if self.rank > 0:
            self.q_proj = _FactorizedLinear(
                self.hidden_size, self.all_head_size, self.rank, bias=True
            )
            self.k_proj = _FactorizedLinear(
                self.hidden_size,
                self.num_kv_heads * self.head_dim,
                self.rank,
                bias=True,
            )
            self.v_proj = _FactorizedLinear(
                self.hidden_size,
                self.num_kv_heads * self.head_dim,
                self.rank,
                bias=True,
            )
            self.out_proj = _FactorizedLinear(
                self.all_head_size, self.hidden_size, self.rank, bias=True
            )
        else:
            self.q_proj = nn.Linear(self.hidden_size, self.all_head_size)
            self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim)
            self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim)
            self.out_proj = nn.Linear(self.all_head_size, self.hidden_size)

        self.dropout = nn.Dropout(dropout)

        if old_attn is not None:
            self.copy_weights(old_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 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 dense(self):
        return self.out_proj

    @property
    def c_proj(self):
        return self.out_proj

    # ------------------------------------------------------------------
    # RE-26: _copy_into for _FactorizedLinear
    # ------------------------------------------------------------------
    def _copy_into(
        self,
        target_proj: nn.Module,
        w: torch.Tensor | None,
        b: torch.Tensor | None,
        target_out: int,
        target_in: int,
    ) -> None:
        if w is None:
            return
        with torch.no_grad():
            if isinstance(target_proj, _FactorizedLinear):
                try:
                    W = w.detach().float()
                    if W.shape[0] != target_out or W.shape[1] != target_in:
                        logger.warning(
                            "Shape mismatch in copy_into: source %s, target (%d, %d).",
                            tuple(W.shape),
                            target_out,
                            target_in,
                        )
                    # Use up to target_proj.rank.
                    req_rank = target_proj.rank
                    if W.shape[0] > 512 and W.shape[1] > 512:
                        U, S, Vh = _randomized_svd(W, req_rank)
                    else:
                        U, S, Vh = torch.linalg.svd(W, full_matrices=False)
                    effective_rank = min(req_rank, U.size(1), Vh.size(0))
                    # RE-26: zero first, then fill only effective_rank rows/cols.
                    target_proj.down.weight.data.zero_()
                    target_proj.up.weight.data.zero_()
                    target_proj.down.weight.data[:effective_rank, :].copy_(
                        Vh[:effective_rank, :].to(target_proj.down.weight.dtype)
                    )
                    US = U[:, :effective_rank] @ torch.diag(S[:effective_rank])
                    target_proj.up.weight.data[:, :effective_rank].copy_(
                        US.to(target_proj.up.weight.dtype)
                    )
                    if b is not None and target_proj.up.bias is not None:
                        n = min(b.numel(), target_proj.up.bias.numel())
                        target_proj.up.bias.data[:n].copy_(b[:n].to(target_proj.up.bias.dtype))
                except Exception as e:
                    logger.error("SVD init failed: %s", e)
            else:
                if w.shape != target_proj.weight.shape:
                    raise ValueError(
                        f"Weight shape mismatch: source {tuple(w.shape)} != "
                        f"target {tuple(target_proj.weight.shape)}."
                    )
                target_proj.weight.data.copy_(w.to(target_proj.weight.dtype))
                if b is not None and target_proj.bias is not None:
                    if b.shape != target_proj.bias.shape:
                        raise ValueError(
                            f"Bias shape mismatch: {tuple(b.shape)} != "
                            f"{tuple(target_proj.bias.shape)}"
                        )
                    target_proj.bias.data.copy_(b.to(target_proj.bias.dtype))

    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        if old_attn is None:
            return

        src = getattr(old_attn, "base_layer", old_attn)
        w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

        # Pattern 1: GPT-2 Conv1D / Linear fused
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn = src.c_attn
            c_proj = src.c_proj
            w_all = c_attn.weight.data
            b_all = c_attn.bias.data if c_attn.bias is not None else None

            hidden = w_all.shape[0] if w_all.shape[0] < w_all.shape[1] else w_all.shape[1]
            if w_all.shape[0] == 3 * hidden or w_all.shape[1] == 3 * hidden:
                if w_all.shape[1] == 3 * hidden:
                    chunks = w_all.chunk(3, dim=1)
                    w_q, w_k, w_v = [c.t().contiguous() for c in chunks]
                else:
                    chunks = w_all.chunk(3, dim=0)
                    w_q, w_k, w_v = chunks[0], chunks[1], chunks[2]
                if b_all is not None:
                    b_q, b_k, b_v = b_all.chunk(3, dim=0)

            w_proj = c_proj.weight.data
            if w_proj.shape[0] != self.hidden_size and w_proj.shape[1] == self.hidden_size:
                w_o = w_proj.t().contiguous()
            else:
                w_o = w_proj
            b_o = c_proj.bias.data if c_proj.bias is not None else None

        # Pattern 2: BERT / RoBERTa nested
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            w_q = s.query.weight.data
            b_q = s.query.bias.data if s.query.bias is not None else None
            w_k = s.key.weight.data
            b_k = s.key.bias.data if s.key.bias is not None else None
            w_v = s.value.weight.data
            b_v = s.value.bias.data if s.value.bias is not None else None
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                w_o = src.output.dense.weight.data
                b_o = src.output.dense.bias.data if src.output.dense.bias is not None else None
            self.is_nested_encoder = True
            self.out_proj = nn.Identity()

        # Pattern 3: Flat query/key/value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.data
            b_q = src.query.bias.data if src.query.bias is not None else None
            w_k = src.key.weight.data
            b_k = src.key.bias.data if src.key.bias is not None else None
            w_v = src.value.weight.data
            b_v = src.value.bias.data if src.value.bias is not None else None
            for attr in ("dense", "out_proj", "o_proj", "proj"):
                if hasattr(src, attr):
                    proj = getattr(src, attr)
                    if hasattr(proj, "weight"):
                        w_o = proj.weight.data
                        b_o = proj.bias.data if getattr(proj, "bias", None) is not None else None
                    break
            if w_o is None:
                self.is_nested_encoder = True
                self.out_proj = nn.Identity()

        # Pattern 4: DistilBERT
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.data
            b_q = src.q_lin.bias.data if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.data
            b_k = src.k_lin.bias.data if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.data
            b_v = src.v_lin.bias.data if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.data
                b_o = src.out_lin.bias.data if src.out_lin.bias is not None else None

        # Pattern 5: Qwen / LLaMA
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.data
            b_q = src.q_proj.bias.data if getattr(src.q_proj, "bias", None) is not None else None
            w_k = src.k_proj.weight.data
            b_k = src.k_proj.bias.data if getattr(src.k_proj, "bias", None) is not None else None
            w_v = src.v_proj.weight.data
            b_v = src.v_proj.bias.data if getattr(src.v_proj, "bias", None) is not None else None
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.data
                b_o = (
                    src.o_proj.bias.data if getattr(src.o_proj, "bias", None) is not None else None
                )

        # Fallback
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                if w_q is None:
                    w_q, b_q = adapter.get_weight_bias(old_attn, "q")
                if w_k is None:
                    w_k, b_k = adapter.get_weight_bias(old_attn, "k")
                if w_v is None:
                    w_v, b_v = adapter.get_weight_bias(old_attn, "v")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(old_attn, "o")
            except Exception as e:
                logger.warning("UniversalAdapter fallback failed: %s", e)

        # RE-28: use class-level flag
        if not ReplaceAttention._first_copy_done:
            for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v), ("w_o", w_o)]:
                if w is not None:
                    norm = float(torch.linalg.norm(w.float()).item())
                    if norm < 1e-4:
                        logger.warning(
                            "Copied weight %s has near-zero Frobenius norm (%.2e)",
                            name,
                            norm,
                        )
            ReplaceAttention._first_copy_done = True

        # RE-27: cast only the affected projections, not the whole module.
        if w_q is not None and hasattr(self.q_proj, "weight"):
            if w_q.dtype != self.q_proj.weight.dtype:
                self.q_proj.to(w_q.dtype)
        if w_k is not None and hasattr(self.k_proj, "weight"):
            if w_k.dtype != self.k_proj.weight.dtype:
                self.k_proj.to(w_k.dtype)
        if w_v is not None and hasattr(self.v_proj, "weight"):
            if w_v.dtype != self.v_proj.weight.dtype:
                self.v_proj.to(w_v.dtype)
        if w_o is not None and not self.is_nested_encoder and hasattr(self.out_proj, "weight"):
            if w_o.dtype != self.out_proj.weight.dtype:
                self.out_proj.to(w_o.dtype)

        self._copy_into(self.q_proj, w_q, b_q, self.all_head_size, self.hidden_size)
        self._copy_into(self.k_proj, w_k, b_k, self.num_kv_heads * self.head_dim, self.hidden_size)
        self._copy_into(self.v_proj, w_v, b_v, self.num_kv_heads * self.head_dim, self.hidden_size)

        if not self.is_nested_encoder:
            self._copy_into(self.out_proj, w_o, b_o, self.hidden_size, self.all_head_size)

    # ------------------------------------------------------------------
    # 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,
        output_attentions: bool = False,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("ReplaceAttention: cross-attention not supported.")

        B, N, _ = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        if self.num_kv_heads < self.num_heads:
            repeat = self.num_heads // self.num_kv_heads
            K = K.repeat_interleave(repeat, dim=1)
            V = V.repeat_interleave(repeat, dim=1)

        is_causal = self.is_decoder and N > 1

        attn_mask_sdpa = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                attn_mask_sdpa = (
                    (attention_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
                ).to(Q.dtype)
            elif attention_mask.dim() == 4:
                attn_mask_sdpa = attention_mask.to(Q.dtype)

        if not output_attentions and head_mask is None:
            try:
                q_sdpa = Q.half() if Q.dtype == torch.float32 else Q
                k_sdpa = K.half() if K.dtype == torch.float32 else K
                v_sdpa = V.half() if V.dtype == torch.float32 else V

                out_sdpa = F.scaled_dot_product_attention(
                    q_sdpa,
                    k_sdpa,
                    v_sdpa,
                    attn_mask=attn_mask_sdpa,
                    is_causal=is_causal and attn_mask_sdpa is None,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                out = out_sdpa.to(dtype).transpose(1, 2).contiguous().view(B, N, self.all_head_size)
                out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

                if self.is_nested_encoder:
                    return out, None
                result = self.out_proj(out)
                result = torch.nan_to_num(result, nan=0.0, posinf=1e4, neginf=-1e4)
                return result, None
            except Exception as e:
                logger.debug("SDPA path failed (%s); falling back.", e)

        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if is_causal:
            causal = (
                torch.triu(
                    torch.full((N, N), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        if attention_mask is not None:
            if attention_mask.dim() == 2:
                scores = scores + (attention_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif attention_mask.dim() == 3:
                scores = scores + attention_mask.unsqueeze(1).float()
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(dtype)
        probs = self.dropout(probs)

        # RE-25: proper head_mask handling.
        if head_mask is not None:
            if head_mask.dim() == 1:
                hm = head_mask.view(1, -1, 1, 1)
            elif head_mask.dim() == 2:
                # (num_layers, num_heads) or (B, H): pick the right slice.
                if head_mask.shape[0] > 1 and head_mask.shape[1] == self.num_heads:
                    hm = head_mask[:1, :].view(1, -1, 1, 1)
                else:
                    hm = head_mask.view(head_mask.shape[0], -1, 1, 1)
            elif head_mask.dim() == 3:
                # Treat as (B, H, N) — mask per head; average over N.
                # Simpler: reduce to (B, H, 1, 1).
                hm = head_mask.mean(dim=-1).view(head_mask.shape[0], -1, 1, 1)
            elif head_mask.dim() == 4:
                hm = head_mask
            else:
                hm = None
            if hm is not None:
                probs = probs * hm.to(probs.dtype)

        context = torch.matmul(probs, V)

        valid_mask = _get_valid_mask(attention_mask, B, N, device)
        context = context * valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)

        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            return context, (probs if output_attentions else None)

        out = self.out_proj(context)
        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

        return out, (probs if output_attentions else None)

Robust identity / low-rank SVD replacement attention.

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.

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

class ResSVDAttention (config: Any,
base_rank: int | None = None,
residual_rank: int | None = None,
rank: int | None = 16,
scaling: float = 1.0,
factorize_base: bool = True,
freeze_base: bool = True,
dropout: float | None = None,
pretrained_attn: nn.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("res_svd")
@ATTENTION_REGISTRY.register("ressvd")
@ATTENTION_REGISTRY.register("res_svd_attention")
@ATTENTION_REGISTRY.register("spectral_residual_svd")
class ResSVDAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any,
        base_rank: int | None = None,
        residual_rank: int | None = None,
        rank: int | None = 16,
        scaling: float = 1.0,
        factorize_base: bool = True,
        freeze_base: bool = True,
        dropout: float | None = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        effective_rank = int(rank) if rank is not None else 16
        # RES-14: pass rank explicitly
        try:
            super().__init__(config, effective_rank)
        except TypeError:
            super().__init__(config)

        comp_config = getattr(config, "compression_config", None) if config else None

        self.hidden_size: int = int(getattr(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.dim = self.hidden_size
        self.num_heads: int = int(
            getattr(config, "num_attention_heads", kwargs.get("num_attention_heads", 12))
        )
        self.num_kv_heads: int = int(
            getattr(
                config,
                "num_key_value_heads",
                kwargs.get("num_key_value_heads", self.num_heads),
            )
        )

        cfg_head_dim = getattr(config, "head_dim", None)
        kw_head_dim = kwargs.get("head_dim")
        inferred = cfg_head_dim or kw_head_dim
        if inferred is None and pretrained_attn is not None:
            for attr in ("q_proj", "query", "q_lin"):
                if hasattr(pretrained_attn, attr):
                    layer_q = getattr(pretrained_attn, attr)
                    if hasattr(layer_q, "weight"):
                        inferred = layer_q.weight.shape[0] // max(1, self.num_heads)
                        break

        self.head_dim = int(inferred or (self.dim // max(1, self.num_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

        self.base_rank: int = int(
            base_rank
            if base_rank is not None
            else getattr(comp_config, "base_rank", effective_rank)
        )
        self.residual_rank: int = int(
            residual_rank
            if residual_rank is not None
            else getattr(comp_config, "residual_rank", max(4, effective_rank // 4))
        )
        self.scaling = float(scaling)
        self.factorize_base = bool(factorize_base)
        self.freeze_base = bool(freeze_base)

        # RES-09: unified decoder detection
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            mt = str(getattr(config, "model_type", "") or "").lower()
            self.is_decoder = (
                bool(getattr(config, "is_decoder", False))
                or bool(getattr(config, "is_causal", False))
                or any(
                    mt.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                or (pretrained_attn is not None and hasattr(pretrained_attn, "c_attn"))
                or bool(kwargs.get("is_causal", False))
                or bool(kwargs.get("is_decoder", False))
            )
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa = True

        # Projections
        self.query = ResSVDLinear(
            in_features=self.hidden_size,
            out_features=self.all_head_size,
            base_rank=self.base_rank,
            residual_rank=self.residual_rank,
            scaling=self.scaling,
            factorize_base=factorize_base,
            bias=True,
        )
        self.key = ResSVDLinear(
            in_features=self.hidden_size,
            out_features=self.all_kv_head_size,
            base_rank=self.base_rank,
            residual_rank=self.residual_rank,
            scaling=self.scaling,
            factorize_base=factorize_base,
            bias=True,
        )
        self.value = ResSVDLinear(
            in_features=self.hidden_size,
            out_features=self.all_kv_head_size,
            base_rank=self.base_rank,
            residual_rank=self.residual_rank,
            scaling=self.scaling,
            factorize_base=factorize_base,
            bias=True,
        )
        self.out_proj = ResSVDLinear(
            in_features=self.all_head_size,
            out_features=self.hidden_size,
            base_rank=self.base_rank,
            residual_rank=self.residual_rank,
            scaling=self.scaling,
            factorize_base=factorize_base,
            bias=True,
        )

        dropout_p = (
            float(dropout)
            if dropout is not None
            else float(
                getattr(
                    config,
                    "dropout",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                )
            )
        )
        self.dropout = nn.Dropout(dropout_p)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

        logger.info(
            "ResSVDAttention: hidden=%d, heads=%d, head_dim=%d, is_decoder=%s, "
            "base_rank=%d, residual_rank=%d",
            self.hidden_size,
            self.num_heads,
            self.head_dim,
            self.is_decoder,
            self.base_rank,
            self.residual_rank,
        )

    # ------------------------------------------------------------------
    @property
    def q_proj(self):
        return self.query

    @property
    def k_proj(self):
        return self.key

    @property
    def v_proj(self):
        return self.value

    @property
    def q_lin(self):
        return self.query

    @property
    def k_lin(self):
        return self.key

    @property
    def v_lin(self):
        return self.value

    @property
    def out_lin(self):
        return self.out_proj

    # ------------------------------------------------------------------
    # RES-02: shape-aware weight transfer via from_linear
    # ------------------------------------------------------------------
    def _init_from_weight(
        self,
        target: ResSVDLinear,
        w: torch.Tensor,
        b: torch.Tensor | None,
    ) -> None:
        """Build a fresh ResSVDLinear from (w, b) and copy its state into target."""
        with torch.no_grad():
            device = target.base_U.device if hasattr(target, "base_U") else torch.device("cpu")
            dtype = target.base_U.dtype if hasattr(target, "base_U") else torch.float32

            # RES-03: temp_linear on same device/dtype
            temp_linear = nn.Linear(
                target.in_features,
                target.out_features,
                bias=(b is not None),
                device=device,
                dtype=dtype,
            )
            # Copy weight/bias into temp_linear
            temp_linear.weight.data.copy_(w.to(device=device, dtype=dtype))
            if b is not None and temp_linear.bias is not None:
                temp_linear.bias.data.copy_(b.to(device=device, dtype=dtype))

            # Build ResSVDLinear and copy state
            new_layer = ResSVDLinear.from_linear(
                temp_linear,
                base_rank=self.base_rank,
                residual_rank=self.residual_rank,
                scaling=self.scaling,
                factorize_base=self.factorize_base,
                freeze_base=self.freeze_base,
            )
            target.load_state_dict(new_layer.state_dict())

    # ------------------------------------------------------------------
    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
        # RES-09: set consistently
        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),
        }

        # 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 hasattr(c_attn_mod, "nf")
                    or hasattr(base_c, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                )
                if is_conv:
                    c_w = c_w.t().contiguous()
                split = 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, :],
                    c_b[:split] if c_b is not None else None,
                )
                extracted["k"] = (
                    c_w[split : 2 * split, :],
                    c_b[split : 2 * split] if c_b is not None else None,
                )
                extracted["v"] = (
                    c_w[2 * split :, :],
                    c_b[2 * split :] 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 hasattr(c_proj_mod, "nf")
                        or hasattr(base_o, "nf")
                        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,
                    )

        # 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,
                )

        # Flat modules
        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 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)
        ]
        # RES-04: raise instead of silently proceeding
        if missing_after:
            raise RuntimeError(
                f"ResSVDAttention.copy_weights: missing {missing_after} "
                f"from {cls_name}. Refusing random init."
            )

        # Handle output projection
        if self.is_nested_encoder and extracted["o"][0] is None:
            # Use identity
            eye = torch.eye(
                self.all_head_size,
                self.hidden_size,
                device=self.out_proj.base_U.device,
                dtype=self.out_proj.base_U.dtype,
            )
            self._init_from_weight(self.out_proj, eye, None)
        elif extracted["o"][0] is not None:
            self._init_from_weight(self.out_proj, extracted["o"][0], extracted["o"][1])

        # q, k, v with GQA resize
        for proj_key, target in (("q", self.query), ("k", self.key), ("v", self.value)):
            w, b = extracted[proj_key]
            if w is None:
                continue

            out_t = target.out_features
            in_t = target.in_features
            src_out, src_in = w.shape

            # Resize for GQA
            if proj_key in ("k", "v"):
                if src_out > out_t and src_out % out_t == 0:
                    rep = src_out // out_t
                    w = w.view(out_t, rep, -1).mean(dim=1)
                    if b is not None:
                        b = b.view(out_t, rep).mean(dim=1)
                elif src_out < out_t and out_t % src_out == 0:
                    rep = out_t // src_out
                    w = w.repeat(rep, 1)
                    if b is not None:
                        b = b.repeat(rep)

            # Shape alignment (pad / truncate)
            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

            self._init_from_weight(target, w, b)

        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)

    # ------------------------------------------------------------------
    # RES-06: robust mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        N: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 3:
            if attention_mask.dtype == torch.bool:
                return attention_mask.squeeze(1)
            if attention_mask.dtype.is_floating_point:
                return (attention_mask > -1e4).squeeze(1)
            return (attention_mask > 0).squeeze(1)
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # 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,
    ):
        # RES-13: explicit error for cross-attention
        if encoder_hidden_states is not None:
            raise NotImplementedError(
                "ResSVDAttention: cross-attention (encoder_hidden_states) is not supported."
            )

        B, N, _ = hidden_states.shape
        H, D = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        Q = self.query(hidden_states).view(B, N, H, D).permute(0, 2, 1, 3)
        K = self.key(hidden_states).view(B, N, self.num_kv_heads, D).permute(0, 2, 1, 3)
        V = self.value(hidden_states).view(B, N, self.num_kv_heads, D).permute(0, 2, 1, 3)

        # RES-05: RoPE with past_len offset
        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]
            cos = cos[..., past_len : past_len + N, :]
            sin = sin[..., past_len : past_len + N, :]
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # GQA expansion
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            K = repeat_kv(K, repeats)
            V = repeat_kv(V, repeats)

        scale = D**-0.5
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        # Causal masking
        is_causal = self.is_decoder and N > 1
        if is_causal:
            causal = (
                torch.triu(
                    torch.full((N, N), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        # Padding mask
        valid_mask = self._prepare_valid_mask(attention_mask, B, N, device)
        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :], -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        # RES-07: unified head_mask
        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)
            if head_mask.size(1) == self.num_heads:
                attn = attn * head_mask.to(attn.dtype)

        context = torch.matmul(attn, V)
        context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)

        return output, (attn if output_attentions else None)

Every 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.

Instance variables

prop k_lin

Expand source code

@property
def k_lin(self):
    return self.key

prop k_proj

Expand source code

@property
def k_proj(self):
    return self.key

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.query

prop q_proj

Expand source code

@property
def q_proj(self):
    return self.query

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.value

prop v_proj

Expand source code

@property
def v_proj(self):
    return self.value

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class ResSVDLinear (in_features: int,
out_features: int,
base_rank: int = 32,
residual_rank: int = 8,
scaling: float = 1.0,
factorize_base: bool = True,
bias: bool = True,
dtype: torch.dtype = torch.float32,
device: torch.device | None = None)

Expand source code

class ResSVDLinear(nn.Module):
    """
    Spectral Residual LoRA (Res-SVD) Linear Layer.

    Executes forward pass as:
    y = x @ W_base^T + scaling * (x @ A^T) @ B^T + bias
    where W_base is either factorized as (base_V, base_U) with requires_grad=False,
    and (lora_A, lora_B) are trainable adapter matrices initialized from the residual tail.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        base_rank: int = 32,
        residual_rank: int = 8,
        scaling: float = 1.0,
        factorize_base: bool = True,
        bias: bool = True,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        max_possible_rank = min(in_features, out_features)

        # Bounds validation
        self.base_rank = max(1, min(base_rank, max_possible_rank))
        remaining_rank = max(0, max_possible_rank - self.base_rank)
        self.residual_rank = max(0, min(residual_rank, remaining_rank))
        self.scaling = float(scaling)
        self.factorize_base = factorize_base

        # Base parameters (Frozen)
        if self.factorize_base:
            self.base_V = nn.Parameter(
                torch.empty(self.base_rank, in_features, dtype=dtype, device=device),
                requires_grad=False,
            )
            self.base_U = nn.Parameter(
                torch.empty(out_features, self.base_rank, dtype=dtype, device=device),
                requires_grad=False,
            )
        else:
            self.base_weight = nn.Parameter(
                torch.empty(out_features, in_features, dtype=dtype, device=device),
                requires_grad=False,
            )

        # Trainable Residual LoRA Adapters
        if self.residual_rank > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.residual_rank, in_features, dtype=dtype, device=device),
                requires_grad=True,
            )
            self.lora_B = nn.Parameter(
                torch.empty(out_features, self.residual_rank, dtype=dtype, device=device),
                requires_grad=True,
            )
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        """Kaiming uniform parameter initialization for base and zero for adapters."""
        if self.factorize_base:
            nn.init.kaiming_uniform_(self.base_V, a=math.sqrt(5))
            nn.init.kaiming_uniform_(self.base_U, a=math.sqrt(5))
        else:
            nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5))

        if self.residual_rank > 0:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

        if self.bias is not None:
            fan_in = self.in_features
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix for compatibility."""
        if self.factorize_base:
            base = torch.matmul(self.base_U, self.base_V)
        else:
            base = self.base_weight
        if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            res = self.scaling * torch.matmul(self.lora_B, self.lora_A)
            return base + res
        return base

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward execution:
        y = Base(x) + scaling * (x @ A^T) @ B^T + bias
        """
        # RESEARCH FIX: Two-stage factorized forward for frozen base to minimize memory footprint
        if self.factorize_base:
            base_intermediate = F.linear(x, self.base_V)
            base_out = F.linear(base_intermediate, self.base_U)
        else:
            base_out = F.linear(x, self.base_weight)

        # Trainable spectral residual LoRA adapter forward
        if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            res_intermediate = F.linear(x, self.lora_A)
            res_out = self.scaling * F.linear(res_intermediate, self.lora_B)
            out = base_out + res_out
        else:
            out = base_out

        if self.bias is not None:
            out = out + self.bias

        return out

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear,
        base_rank: int = 32,
        residual_rank: int = 8,
        scaling: float = 1.0,
        factorize_base: bool = True,
        freeze_base: bool = True,
    ) -> ResSVDLinear:
        """
        Construct ResSVDLinear by decomposing a pretrained nn.Linear weight:
        1. Base factors U_r, V_r from the dominant top-r singular components.
        2. Residual adapter factors A, B directly initialized from the residual tail (r+1 to r+k).

        Args:
            linear: Source dense linear layer.
            base_rank: Rank r for the frozen base matrix.
            residual_rank: Rank k for the trainable residual adapter.
            scaling: Multiplier alpha / k for adapter output.
            factorize_base: Whether to store base weights as factorized U, V matrices.
            freeze_base: Whether to freeze base factors (requires_grad=False).

        Returns:
            ResSVDLinear: Spectral residual LoRA module.
        """
        in_f = linear.in_features
        out_f = linear.out_features
        max_rank = min(in_f, out_f)

        r = max(1, min(base_rank, max_rank))
        rem = max(0, max_rank - r)
        k = max(0, min(residual_rank, rem))

        new_layer = cls(
            in_features=in_f,
            out_features=out_f,
            base_rank=r,
            residual_rank=k,
            scaling=scaling,
            factorize_base=factorize_base,
            bias=linear.bias is not None,
            dtype=linear.weight.dtype,
            device=linear.weight.device,
        )

        W = linear.weight.data.float()  # (out_f, in_f)

        # SVD decomposition of source weight
        U, S, Vh = torch.linalg.svd(W, full_matrices=False)

        # 1. Base Truncation (Dominant Top-r Singular Components)
        U_base = U[:, :r] * torch.sqrt(S[:r]).unsqueeze(0)  # (out_f, r)
        V_base = torch.sqrt(S[:r]).unsqueeze(1) * Vh[:r, :]  # (r, in_f)

        if factorize_base:
            new_layer.base_U.data.copy_(U_base.to(dtype=linear.weight.dtype))
            new_layer.base_V.data.copy_(V_base.to(dtype=linear.weight.dtype))
            new_layer.base_U.requires_grad = not freeze_base
            new_layer.base_V.requires_grad = not freeze_base
        else:
            W_base = (U[:, :r] * S[:r].unsqueeze(0)) @ Vh[:r, :]
            new_layer.base_weight.data.copy_(W_base.to(dtype=linear.weight.dtype))
            new_layer.base_weight.requires_grad = not freeze_base

        # 2. Residual Tail Initialization (Singular directions r to r+k)
        # RESEARCH FIX: Analytical spectral initialization of LoRA adapter weights (A, B) directly
        # from the residual tail E = W - U_r V_r minimizes initial reconstruction loss to zero-shot optimality.
        if k > 0:
            S_tail = S[r : r + k]
            U_tail = U[:, r : r + k]
            Vh_tail = Vh[r : r + k, :]

            # Scale S_tail by adapter scaling factor gamma: (B @ A) * gamma = E_tail
            scale_factor = max(1e-8, float(scaling))
            S_tail_scaled = S_tail / scale_factor

            sqrt_S = torch.sqrt(torch.clamp(S_tail_scaled, min=0.0))
            lora_A_init = sqrt_S.unsqueeze(1) * Vh_tail  # (k, in_f)
            lora_B_init = U_tail * sqrt_S.unsqueeze(0)  # (out_f, k)

            new_layer.lora_A.data.copy_(lora_A_init.to(dtype=linear.weight.dtype))
            new_layer.lora_B.data.copy_(lora_B_init.to(dtype=linear.weight.dtype))
            new_layer.lora_A.requires_grad = True
            new_layer.lora_B.requires_grad = True

        # Copy bias if present
        if linear.bias is not None:
            new_layer.bias.data.copy_(linear.bias.data)

        return new_layer

    def merge_weights(self) -> nn.Linear:
        """
        Merge frozen base weights and trainable residual adapter weights into a dense nn.Linear.

        Returns:
            nn.Linear: Standard dense linear layer with merged weights.
        """
        device = next(self.parameters()).device
        dtype = self.base_U.dtype if self.factorize_base else self.base_weight.dtype

        merged_linear = nn.Linear(
            self.in_features,
            self.out_features,
            bias=self.bias is not None,
            device=device,
            dtype=dtype,
        )

        with torch.no_grad():
            if self.factorize_base:
                W_base = self.base_U.float() @ self.base_V.float()
            else:
                W_base = self.base_weight.float()

            if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
                delta_W = self.scaling * (self.lora_B.float() @ self.lora_A.float())
                W_total = W_base + delta_W
            else:
                W_total = W_base

            merged_linear.weight.data.copy_(W_total.to(dtype=dtype))
            if self.bias is not None:
                merged_linear.bias.data.copy_(self.bias.data)

        return merged_linear

    def get_residual_error_norm(self, original_weight: torch.Tensor | None = None) -> float:
        """
        Calculate relative Frobenius norm reconstruction error:
        || W_orig - (W_base + scaling * B @ A) ||_F / || W_orig ||_F
        """
        with torch.no_grad():
            if self.factorize_base:
                W_approx = self.base_U.float() @ self.base_V.float()
            else:
                W_approx = self.base_weight.float()

            if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
                W_approx = W_approx + self.scaling * (self.lora_B.float() @ self.lora_A.float())

            if original_weight is not None:
                W_orig = original_weight.float().to(W_approx.device)
                error = torch.norm(W_orig - W_approx, p="fro") / (
                    torch.norm(W_orig, p="fro") + 1e-12
                )
                return float(error.item())

            return 0.0

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"base_rank={self.base_rank}, residual_rank={self.residual_rank}, "
            f"scaling={self.scaling}, factorize_base={self.factorize_base}"
        )

Spectral Residual LoRA (Res-SVD) Linear Layer.

Executes forward pass as: y = x @ W_base^T + scaling * (x @ A^T) @ B^T + bias where W_base is either factorized as (base_V, base_U) with requires_grad=False, and (lora_A, lora_B) are trainable adapter matrices initialized from the residual tail.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Static methods

def from_linear(linear: nn.Linear,
base_rank: int = 32,
residual_rank: int = 8,
scaling: float = 1.0,
factorize_base: bool = True,
freeze_base: bool = True) ‑> ResSVDLinear

Construct ResSVDLinear by decomposing a pretrained nn.Linear weight: 1. Base factors U_r, V_r from the dominant top-r singular components. 2. Residual adapter factors A, B directly initialized from the residual tail (r+1 to r+k).

Args

-----=

linear

Source dense linear layer.

base_rank

Rank r for the frozen base matrix.

residual_rank

Rank k for the trainable residual adapter.

scaling

Multiplier alpha / k for adapter output.

factorize_base

Whether to store base weights as factorized U, V matrices.

freeze_base

Whether to freeze base factors (requires_grad=False).

Returns

-----=

ResSVDLinear

Spectral residual LoRA module.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix for compatibility."""
    if self.factorize_base:
        base = torch.matmul(self.base_U, self.base_V)
    else:
        base = self.base_weight
    if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        res = self.scaling * torch.matmul(self.lora_B, self.lora_A)
        return base + res
    return base

Reconstructed effective weight matrix for compatibility.

Methods

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    return (
        f"in_features={self.in_features}, out_features={self.out_features}, "
        f"base_rank={self.base_rank}, residual_rank={self.residual_rank}, "
        f"scaling={self.scaling}, factorize_base={self.factorize_base}"
    )

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward execution:
    y = Base(x) + scaling * (x @ A^T) @ B^T + bias
    """
    # RESEARCH FIX: Two-stage factorized forward for frozen base to minimize memory footprint
    if self.factorize_base:
        base_intermediate = F.linear(x, self.base_V)
        base_out = F.linear(base_intermediate, self.base_U)
    else:
        base_out = F.linear(x, self.base_weight)

    # Trainable spectral residual LoRA adapter forward
    if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        res_intermediate = F.linear(x, self.lora_A)
        res_out = self.scaling * F.linear(res_intermediate, self.lora_B)
        out = base_out + res_out
    else:
        out = base_out

    if self.bias is not None:
        out = out + self.bias

    return out

Forward execution: y = Base(x) + scaling * (x @ A^T) @ B^T + bias

def get_residual_error_norm(self, original_weight: torch.Tensor | None = None) ‑> float

Expand source code

def get_residual_error_norm(self, original_weight: torch.Tensor | None = None) -> float:
    """
    Calculate relative Frobenius norm reconstruction error:
    || W_orig - (W_base + scaling * B @ A) ||_F / || W_orig ||_F
    """
    with torch.no_grad():
        if self.factorize_base:
            W_approx = self.base_U.float() @ self.base_V.float()
        else:
            W_approx = self.base_weight.float()

        if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            W_approx = W_approx + self.scaling * (self.lora_B.float() @ self.lora_A.float())

        if original_weight is not None:
            W_orig = original_weight.float().to(W_approx.device)
            error = torch.norm(W_orig - W_approx, p="fro") / (
                torch.norm(W_orig, p="fro") + 1e-12
            )
            return float(error.item())

        return 0.0

Calculate relative Frobenius norm reconstruction error: || W_orig - (W_base + scaling * B @ A) ||_F / || W_orig ||_F

def merge_weights(self) ‑> torch.nn.modules.linear.Linear

Expand source code

def merge_weights(self) -> nn.Linear:
    """
    Merge frozen base weights and trainable residual adapter weights into a dense nn.Linear.

    Returns:
        nn.Linear: Standard dense linear layer with merged weights.
    """
    device = next(self.parameters()).device
    dtype = self.base_U.dtype if self.factorize_base else self.base_weight.dtype

    merged_linear = nn.Linear(
        self.in_features,
        self.out_features,
        bias=self.bias is not None,
        device=device,
        dtype=dtype,
    )

    with torch.no_grad():
        if self.factorize_base:
            W_base = self.base_U.float() @ self.base_V.float()
        else:
            W_base = self.base_weight.float()

        if self.residual_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            delta_W = self.scaling * (self.lora_B.float() @ self.lora_A.float())
            W_total = W_base + delta_W
        else:
            W_total = W_base

        merged_linear.weight.data.copy_(W_total.to(dtype=dtype))
        if self.bias is not None:
            merged_linear.bias.data.copy_(self.bias.data)

    return merged_linear

Merge frozen base weights and trainable residual adapter weights into a dense nn.Linear.

Returns

-----=

nn.Linear

Standard dense linear layer with merged weights.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    """Kaiming uniform parameter initialization for base and zero for adapters."""
    if self.factorize_base:
        nn.init.kaiming_uniform_(self.base_V, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.base_U, a=math.sqrt(5))
    else:
        nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5))

    if self.residual_rank > 0:
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

    if self.bias is not None:
        fan_in = self.in_features
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        nn.init.uniform_(self.bias, -bound, bound)

Kaiming uniform parameter initialization for base and zero for adapters.

class RoutingAttention (config: Any = None,
rank: int | None = None,
num_clusters: int = 16,
n_iters: int = 3,
pretrained_attn: nn.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("routing_class")
@ATTENTION_REGISTRY.register("routing_attention_class")
class RoutingAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any = None,
        rank: int | None = None,
        num_clusters: int = 16,
        n_iters: int = 3,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank or 64)
        except TypeError:
            super().__init__(config)

        comp_config = getattr(config, "compression_config", None) if config else None

        self.dim = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.num_heads = int(
            _cfg_get(config, "num_attention_heads", kwargs.get("num_attention_heads", 12))
        )
        self.num_kv_heads = int(
            _cfg_get(
                config,
                "num_key_value_heads",
                kwargs.get("num_key_value_heads", self.num_heads),
            )
        )

        if self.num_heads <= 0 or self.dim % self.num_heads != 0:
            raise ValueError(f"dim={self.dim} not divisible by num_heads={self.num_heads}")

        self.head_dim = self.dim // self.num_heads
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim
        self.num_clusters = int(
            _cfg_get(
                comp_config,
                "num_clusters",
                _cfg_get(config, "routing_num_clusters", num_clusters),
            )
        )
        self.n_iters = max(1, min(int(n_iters), 10))

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            mt = str(_cfg_get(config, "model_type", "") or "").lower()
            self.is_decoder = (
                bool(_cfg_get(config, "is_decoder", False))
                or bool(kwargs.get("is_causal", False))
                or any(
                    mt.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                or (pretrained_attn is not None and hasattr(pretrained_attn, "c_attn"))
            )
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        self.lora_rank = int(
            _cfg_get(
                config,
                "lora_rank",
                _cfg_get(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
        )
        self.lora_alpha = float(
            _cfg_get(
                config,
                "lora_alpha",
                _cfg_get(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
        )
        self.proj_rank = int(
            kwargs.get(
                "proj_rank",
                _cfg_get(comp_config, "proj_rank", _cfg_get(config, "proj_rank", 64)),
            )
        )

        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa = True

        # Factorized projections
        FF = FactorizedLoRALinear
        self.q_proj = FF(
            self.dim,
            self.all_head_size,
            rank=self.proj_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            bias=True,
        )
        self.k_proj = 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.v_proj = 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,
        )

        dropout_p = _cfg_get(
            config, "dropout", _cfg_get(config, "attention_probs_dropout_prob", 0.1)
        )
        self.dropout = nn.Dropout(dropout_p)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    # ------------------------------------------------------------------
    @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

    # ------------------------------------------------------------------
    @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("RoutingAttention._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_()

    # ------------------------------------------------------------------
    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),
        }

        # 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[0] < c_w.shape[1])
                )
                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,
                    )

        # Nested / flat
        elif hasattr(src, "self"):
            s = src
            while hasattr(s, "self") and not any(
                hasattr(s, a) for a in ("query", "key", "value", "q_proj", "k_proj", "v_proj")
            ):
                s = s.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,
                )
        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 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"RoutingAttention.copy_weights: missing {missing_after} from {cls_name}."
            )

        for proj_key, target in (
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ):
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    eye = torch.eye(
                        self.all_head_size,
                        self.all_head_size,
                        device=target.down.weight.device,
                        dtype=target.down.weight.dtype,
                    )
                    self._svd_init(eye, None, target)
                continue

            # GQA resize for k, v
            if proj_key in ("k", "v"):
                out_t = target.out_features
                src_out = w.shape[0]
                if src_out > out_t and src_out % out_t == 0:
                    rep = src_out // out_t
                    w = w.view(out_t, rep, -1).mean(dim=1)
                    if b is not None:
                        b = b.view(out_t, rep).mean(dim=1)
                elif src_out < out_t and out_t % src_out == 0:
                    rep = out_t // src_out
                    w = w.repeat(rep, 1)
                    if b is not None:
                        b = b.repeat(rep)

            self._svd_init(w, b, target)
            target.freeze_base_weights()

        if hasattr(src, "is_decoder"):
            self.is_decoder = bool(src.is_decoder)

    # ------------------------------------------------------------------
    # Robust valid-mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(
        attention_mask: torch.Tensor | None,
        B: int,
        L: int,
        device: torch.device,
    ) -> torch.Tensor:
        if attention_mask is None:
            return torch.ones(B, L, device=device, dtype=torch.bool)
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 3:
            if attention_mask.dtype == torch.bool:
                return attention_mask.squeeze(1)
            if attention_mask.dtype.is_floating_point:
                return (attention_mask > -1e4).squeeze(1)
            return (attention_mask > 0).squeeze(1)
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # Stable k-means
    # ------------------------------------------------------------------
    @staticmethod
    def _kmeans_assign(feat: torch.Tensor, K: int, n_iters: int) -> torch.Tensor:
        """
        Deterministic k-means assignment on (B, H, L, d).
        Uses evenly-spaced initialization + squared-distance expansion.
        """
        B, H, L, d = feat.shape
        K = min(K, L)
        step = max(1, L // K)
        init_idx = torch.arange(0, L, step, device=feat.device)[:K]
        if init_idx.size(0) < K:
            extra = torch.arange(0, K - init_idx.size(0), device=feat.device)
            init_idx = torch.cat([init_idx, extra])

        feat_f = feat.float()
        centroids = feat_f[:, :, init_idx, :].clone()  # (B, H, K, d)
        cluster_ids = torch.zeros(B, H, L, dtype=torch.long, device=feat.device)

        for _ in range(n_iters):
            # (B, H, L, K) via squared-distance expansion
            x_norm = (feat_f**2).sum(dim=-1, keepdim=True)  # (B, H, L, 1)
            c_norm = (centroids**2).sum(dim=-1).unsqueeze(-2)  # (B, H, 1, K)
            cross = torch.matmul(feat_f, centroids.transpose(-2, -1))  # (B, H, L, K)
            dist = x_norm - 2 * cross + c_norm
            cluster_ids = torch.argmin(dist, dim=-1)

            one_hot = F.one_hot(cluster_ids, num_classes=K).to(feat_f.dtype)
            counts = one_hot.sum(dim=2).clamp(min=1.0).unsqueeze(-1)
            sum_f = torch.einsum("bhlk,bhld->bhkd", one_hot, feat_f)
            centroids = sum_f / counts

        return cluster_ids

    # ------------------------------------------------------------------
    # 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,
        output_attentions: bool | None = False,
        layer_idx: int = 0,
        *args: Any,
        **kwargs: Any,
    ):
        B, L, D = hidden_states.shape
        H, d = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        # --- Projections ---
        Q = self.q_proj(hidden_states).view(B, L, H, d).permute(0, 2, 1, 3).contiguous()
        K = (
            self.k_proj(hidden_states)
            .view(B, L, self.num_kv_heads, d)
            .permute(0, 2, 1, 3)
            .contiguous()
        )
        V = (
            self.v_proj(hidden_states)
            .view(B, L, self.num_kv_heads, d)
            .permute(0, 2, 1, 3)
            .contiguous()
        )

        # --- RoPE ---
        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]
            cos = cos[..., past_len : past_len + L, :]
            sin = sin[..., past_len : past_len + L, :]
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # --- GQA ---
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            K = repeat_kv(K, repeats)
            V = repeat_kv(V, repeats)

        # --- Valid mask ---
        valid_mask_2d = self._prepare_valid_mask(attention_mask, B, L, device)  # (B, L) bool

        # --- Clustering ---
        K_num = min(max(1, self.num_clusters), L)
        joint_feat = F.normalize(Q + K, p=2, dim=-1)  # (B, H, L, d)
        cluster_ids = self._kmeans_assign(joint_feat, K_num, self.n_iters)  # (B, H, L) long

        scale = d**-0.5
        is_causal = self.is_decoder or kwargs.get("is_causal", False)

        # ===========================================================
        # Sort + Gather + Chunk attention (real O(N²/K))
        # ===========================================================
        # For each cluster k:
        #   1. Use sort-trick to gather cluster members into a compact block
        #   2. Compute attention on the small (max_q × max_k) matrix
        #   3. Scatter the result back to the original positions
        # ===========================================================

        outputs = torch.zeros_like(Q)  # (B, H, L, d)
        arange_L = torch.arange(L, device=device)

        for k in range(K_num):
            q_mask_k = cluster_ids == k  # (B, H, L) bool
            # Keys: cluster members + token 0 (global anchor) ∩ valid
            k_mask_k = q_mask_k | (arange_L.view(1, 1, -1) == 0)
            k_mask_k = k_mask_k & valid_mask_2d.unsqueeze(1)  # (B, H, L)

            # Determine compact block sizes
            max_q = int(q_mask_k.sum(dim=-1).max().item())
            max_k = int(k_mask_k.sum(dim=-1).max().item())

            if max_q == 0 or max_k == 0:
                continue

            # Sort trick: cluster members get score 2L - pos, others get -pos
            # → argsort descending puts cluster members first (ordered by position desc)
            q_score = q_mask_k.to(torch.int64) * (2 * L) - arange_L.view(1, 1, L)
            k_score = k_mask_k.to(torch.int64) * (2 * L) - arange_L.view(1, 1, L)

            q_idx = q_score.argsort(dim=-1, descending=True)[:, :, :max_q]  # (B, H, max_q)
            k_idx = k_score.argsort(dim=-1, descending=True)[:, :, :max_k]  # (B, H, max_k)

            # Padding mask (which entries are real cluster members)
            valid_q = torch.gather(q_mask_k, 2, q_idx)  # (B, H, max_q)
            valid_k = torch.gather(k_mask_k, 2, k_idx)  # (B, H, max_k)

            # Gather Q, K, V for this cluster
            q_idx_exp = q_idx.unsqueeze(-1).expand(-1, -1, -1, d)
            k_idx_exp = k_idx.unsqueeze(-1).expand(-1, -1, -1, d)

            Q_k = torch.gather(Q, 2, q_idx_exp)  # (B, H, max_q, d)
            K_k = torch.gather(K, 2, k_idx_exp)  # (B, H, max_k, d)
            V_k = torch.gather(V, 2, k_idx_exp)  # (B, H, max_k, d)

            # Compact attention: (B, H, max_q, max_k) — this is O((N/K)²), not O(N²)
            scores_k = torch.matmul(Q_k.float(), K_k.float().transpose(-2, -1).contiguous()) * scale

            # Build pair mask
            pair_mask = valid_q.unsqueeze(-1) & valid_k.unsqueeze(-2)

            # Causal mask via ABSOLUTE positions (preserves causality after sort)
            if is_causal:
                abs_pos_q = q_idx.unsqueeze(-1)  # (B, H, max_q, 1)
                abs_pos_k = k_idx.unsqueeze(-2)  # (B, H, 1, max_k)
                causal_mask = abs_pos_q >= abs_pos_k
                pair_mask = pair_mask & causal_mask

            scores_k = scores_k.masked_fill(~pair_mask, -1e4)
            scores_k = torch.nan_to_num(scores_k, nan=-1e4, posinf=1e4, neginf=-1e4)

            attn_k = F.softmax(scores_k, dim=-1).to(dtype)
            attn_k = torch.nan_to_num(attn_k, nan=0.0)
            attn_k = self.dropout(attn_k)

            out_k = torch.matmul(attn_k, V_k)  # (B, H, max_q, d)

            # Zero out padded queries before scattering
            out_k = out_k * valid_q.unsqueeze(-1).to(out_k.dtype)

            # Scatter back (each token belongs to exactly one cluster → safe)
            outputs.scatter_add_(2, q_idx_exp, out_k.to(outputs.dtype))

        # ===========================================================
        # Post-processing
        # ===========================================================

        # 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
            outputs = outputs * hm.to(outputs.dtype)

        # Zero out invalid query positions
        outputs = outputs * valid_mask_2d.unsqueeze(1).unsqueeze(-1).to(outputs.dtype)

        # Reshape
        context = outputs.permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            return context, None

        output = self.out_proj(context)
        output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4)
        return output, None

Every 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.

Instance variables

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 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

class SALTAttention (config: Any,
rank: int | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
is_decoder: bool | None = None,
use_fisher: bool | None = None,
adaptive_rank: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("salt_f_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("salt_f", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("salt_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("salt", supports_causal=True, supports_gqa=True)
class SALTAttention(BaseAttention):
    r"""
    SALT Attention module replacing standard linear projections with SALTLinear.
    Hardened for production across DistilBERT, RoBERTa, GPT-2, and Qwen3.5.
    Supports Fisher-weighted SVD, attention-aware gamma, and adaptive rank (SALT-F).
    """

    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        is_decoder: bool | None = None,
        use_fisher: bool | None = None,
        adaptive_rank: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        self.num_heads: int = int(
            _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "num_heads", kwargs.get("num_attention_heads", 4)),
            )
        )
        self.hidden_size: int = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.head_dim: int = int(
            _cfg_get(
                config,
                "head_dim",
                self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
            )
        )
        self.all_head_size: int = self.num_heads * self.head_dim

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = max(1, self.num_heads // self.num_kv_heads)
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        self.rank: int = (
            rank
            if rank is not None
            else int(
                _cfg_get(
                    config,
                    "salt_rank",
                    _cfg_get(config, "rank", kwargs.pop("rank", 16)),
                )
            )
        )

        # Adaptive rank & Fisher flags
        if adaptive_rank is not None:
            self.adaptive_rank = bool(adaptive_rank)
        else:
            self.adaptive_rank = bool(
                _cfg_get(
                    config,
                    "adaptive_rank",
                    _cfg_get(config, "salt_adaptive_rank", kwargs.get("adaptive_rank", True)),
                )
            )

        if use_fisher is not None:
            self.use_fisher = bool(use_fisher)
        else:
            self.use_fisher = bool(
                _cfg_get(
                    config,
                    "use_fisher",
                    _cfg_get(config, "salt_use_fisher", kwargs.get("use_fisher", True)),
                )
            )

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(_cfg_get(config, "model_type", "") or "").lower()
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            self.is_decoder = bool(
                _cfg_get(config, "is_decoder", False)
                or _cfg_get(config, "causal", False)
                or any(_mt.startswith(p) for p in _dec_pfx)
            )

        self.is_nested_encoder: bool = False

        dropout_p = getattr(
            config,
            "dropout",
            getattr(
                config,
                "hidden_dropout_prob",
                getattr(config, "attention_probs_dropout_prob", 0.1),
            ),
        )
        self.dropout = nn.Dropout(dropout_p)

        self.proj_ranks = self._get_projection_ranks(pretrained_attn, self.rank)
        self.q_proj = SALTLinear(
            self.hidden_size,
            self.all_head_size,
            self.proj_ranks["q"],
            use_fisher=self.use_fisher,
        )
        self.k_proj = SALTLinear(
            self.hidden_size,
            self.kv_head_size,
            self.proj_ranks["k"],
            use_fisher=self.use_fisher,
        )
        self.v_proj = SALTLinear(
            self.hidden_size,
            self.kv_head_size,
            self.proj_ranks["v"],
            use_fisher=self.use_fisher,
        )
        self.out_proj = SALTLinear(
            self.all_head_size,
            self.hidden_size,
            self.proj_ranks["o"],
            use_fisher=self.use_fisher,
        )

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    def _get_projection_ranks(
        self, pretrained_attn: nn.Module | None, base_rank: int
    ) -> dict[str, int]:
        """
        Compute per-projection rank based on SVD energy.
        Returns dict {'q': r_q, 'k': r_k, 'v': r_v, 'o': r_o}.
        """
        ranks = {"q": base_rank, "k": base_rank, "v": base_rank, "o": base_rank}
        if not getattr(self, "adaptive_rank", True):
            return ranks

        if pretrained_attn is None:
            return ranks

        try:
            for proj_key, attr_list in [
                ("q", ["q_proj", "q_lin", "query"]),
                ("k", ["k_proj", "k_lin", "key"]),
                ("v", ["v_proj", "v_lin", "value"]),
                ("o", ["o_proj", "out_proj", "out_lin", "dense"]),
            ]:
                for attr in attr_list:
                    if hasattr(pretrained_attn, attr):
                        layer = getattr(pretrained_attn, attr)
                        target = getattr(layer, "base_layer", layer)
                        if hasattr(target, "weight") and target.weight is not None:
                            W = target.weight.data.float()
                            _, S, _ = torch.linalg.svd(W, full_matrices=False)
                            ranks[proj_key] = _compute_energy_rank(
                                S, 0.95, 8, max(base_rank * 2, 32)
                            )
                        break
        except Exception as e:
            logger.debug("Per-projection rank computation failed: %s", e)

        return ranks

    @property
    def query(self) -> SALTLinear:
        return self.q_proj

    @property
    def key(self) -> SALTLinear:
        return self.k_proj

    @property
    def value(self) -> SALTLinear:
        return self.v_proj

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # 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
        ):
            weight = src.c_attn.weight.data
            bias = getattr(src.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                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:
            w_o = src.c_proj.weight.data
            if (
                w_o.dim() == 2
                and w_o.shape[0] != self.hidden_size
                and w_o.shape[1] == self.hidden_size
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(src.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(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"],
            "k": ["k_proj", "k_lin", "key"],
            "v": ["v_proj", "v_lin", "value"],
            "o": ["o_proj", "out_proj", "out_lin", "dense"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ) and not has_out_proj

        # Check required weights
        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."
            )

        # Frobenius norm validation
        for proj_key in ["q", "k", "v"]:
            w = extracted[proj_key][0]
            if w is not None:
                norm = torch.linalg.norm(w.float())
                if norm < 1e-4:
                    logger.warning(
                        f"SALTAttention extracted weight for {proj_key} has near-zero Frobenius norm ({norm.item():.6f})"
                    )

        if extracted["o"][0] is not None:
            norm_o = torch.linalg.norm(extracted["o"][0].float())
            if norm_o < 1e-4:
                logger.warning(
                    f"SALTAttention extracted weight for output projection has near-zero Frobenius norm ({norm_o.item():.6f})"
                )

        def _orient(w: torch.Tensor, in_dim: int, out_dim: int) -> torch.Tensor:
            """Handle both nn.Linear [out, in] and Conv1D [in, out]."""
            if w.dim() != 2:
                return w
            if w.shape == (out_dim, in_dim):
                return w
            if w.shape == (in_dim, out_dim):
                return w.t().contiguous()
            # fallback
            if w.shape[1] == in_dim:
                return w
            if w.shape[0] == in_dim:
                return w.t().contiguous()
            return w

        # Recalculate adaptive ranks from actual extracted weights if available
        if getattr(self, "adaptive_rank", True):
            for proj_key in ["q", "k", "v", "o"]:
                w_cand = extracted[proj_key][0]
                if w_cand is not None:
                    try:
                        _, S_cand, _ = torch.linalg.svd(w_cand.float(), full_matrices=False)
                        self.proj_ranks[proj_key] = _compute_energy_rank(
                            S_cand, 0.95, 8, max(self.rank * 2, 32)
                        )
                    except Exception:
                        pass

        # Initialize SALTLinear projections
        with torch.no_grad():
            for proj_key, attr, in_dim, out_dim in [
                ("q", "q_proj", self.hidden_size, self.all_head_size),
                ("k", "k_proj", self.hidden_size, self.kv_head_size),
                ("v", "v_proj", self.hidden_size, self.kv_head_size),
                ("o", "out_proj", self.all_head_size, self.hidden_size),
            ]:
                w, b = extracted[proj_key]
                proj_r = self.proj_ranks.get(proj_key, self.rank)
                if w is not None:
                    w_oriented = _orient(w, in_dim, out_dim)
                    src_out, src_in = w_oriented.shape

                    # Handle KV dimension matching for GQA
                    if proj_key in ("k", "v"):
                        if src_out == out_dim:
                            final_w = w_oriented
                            final_b = b
                        elif src_out > out_dim and src_out % out_dim == 0:
                            rep = src_out // out_dim
                            final_w = w_oriented.view(out_dim, rep, -1).mean(dim=1)
                            final_b = b.view(out_dim, rep).mean(dim=1) if b is not None else None
                        elif src_out < out_dim and out_dim % src_out == 0:
                            rep = out_dim // src_out
                            final_w = w_oriented.repeat(rep, 1)
                            final_b = b.repeat(rep) if b is not None else None
                        else:
                            final_w = torch.zeros(out_dim, in_dim, device=w.device, dtype=w.dtype)
                            final_w[: min(out_dim, src_out), : min(in_dim, src_in)] = w_oriented[
                                : min(out_dim, src_out), : min(in_dim, src_in)
                            ]
                            final_b = b[: min(out_dim, len(b))] if b is not None else None
                    else:
                        if src_out == out_dim and src_in == in_dim:
                            final_w = w_oriented
                            final_b = b
                        else:
                            final_w = torch.zeros(out_dim, in_dim, device=w.device, dtype=w.dtype)
                            final_w[: min(out_dim, src_out), : min(in_dim, src_in)] = w_oriented[
                                : min(out_dim, src_out), : min(in_dim, src_in)
                            ]
                            final_b = b[: min(out_dim, len(b))] if b is not None else None

                    new_linear = SALTLinear(
                        in_dim,
                        out_dim,
                        proj_r,
                        original_weight=final_w,
                        original_bias=final_b,
                        use_fisher=self.use_fisher,
                    )
                    new_linear.to(device=w.device, dtype=w.dtype)
                    setattr(self, attr, new_linear)
                elif proj_key == "o" and self.is_nested_encoder:
                    eye_w = torch.eye(
                        out_dim,
                        in_dim,
                        device=self.out_proj.U.device,
                        dtype=self.out_proj.U.dtype,
                    )
                    new_linear = SALTLinear(
                        in_dim,
                        out_dim,
                        proj_r,
                        original_weight=eye_w,
                        original_bias=None,
                        use_fisher=self.use_fisher,
                    )
                    setattr(self, attr, new_linear)

        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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, L, D = hidden_states.shape
        H, Dh = self.num_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        L_kv = kv_input.shape[1]

        Q = self.q_proj(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
        K = self.k_proj(kv_input).view(B, L_kv, self.num_kv_heads, Dh).transpose(1, 2).contiguous()
        V = self.v_proj(kv_input).view(B, L_kv, self.num_kv_heads, Dh).transpose(1, 2).contiguous()

        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
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = Dh**-0.5
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_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 = getattr(self, "is_decoder", False) and not is_cross and L > 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, L, self.all_head_size)
                )
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                out = (
                    _sdpa_out
                    if getattr(self, "is_nested_encoder", False)
                    else self.out_proj(_sdpa_out)
                )
                if out.shape[-1] != hidden_states.shape[-1]:
                    out = (
                        F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
                        if out.shape[-1] < hidden_states.shape[-1]
                        else out[..., : hidden_states.shape[-1]]
                    )
                return out, None
            except Exception:
                pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross and L > 1:
            causal_mask = (
                torch.triu(
                    torch.full((L, L_kv), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(dtype)
        attn = self.dropout(attn)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn = attn * head_mask

        context = torch.matmul(attn, V)
        context = context.transpose(1, 2).contiguous().view(B, L, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            out = context
        else:
            out = self.out_proj(context)

        if out.shape[-1] != hidden_states.shape[-1]:
            out = (
                F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
                if out.shape[-1] < hidden_states.shape[-1]
                else out[..., : hidden_states.shape[-1]]
            )

        return out, (attn if output_attentions else None)

SALT Attention module replacing standard linear projections with SALTLinear. Hardened for production across DistilBERT, RoBERTa, GPT-2, and Qwen3.5. Supports Fisher-weighted SVD, attention-aware gamma, and adaptive rank (SALT-F).

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.

Instance variables

prop key : SALTLinear

Expand source code

@property
def key(self) -> SALTLinear:
    return self.k_proj

prop query : SALTLinear

Expand source code

@property
def query(self) -> SALTLinear:
    return self.q_proj

prop value : SALTLinear

Expand source code

@property
def value(self) -> SALTLinear:
    return self.v_proj

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class SALTLinear (in_features: int,
out_features: int,
rank: int = 16,
original_weight: torch.Tensor | None = None,
original_bias: torch.Tensor | None = None,
fisher_diag: torch.Tensor | None = None,
use_fisher: bool = True)

Expand source code

class SALTLinear(nn.Module):
    is_custom_attention: bool = True
    r"""
    Linear projection layer with Fisher-weighted SVD initialization,
    attention-aware singular value modulation (\gamma), and low-rank residual branch (A B).

    Formulation:
        $$W = U_{:, :r} \operatorname{diag}(\Sigma_{1:r} \odot \gamma) V_{1:r, :}^\top + A B$$
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 16,
        original_weight: torch.Tensor | None = None,
        original_bias: torch.Tensor | None = None,
        fisher_diag: torch.Tensor | None = None,
        use_fisher: bool = True,
    ) -> None:
        super().__init__()
        self.in_features: int = in_features
        self.out_features: int = out_features
        self.rank: int = min(rank, min(in_features, out_features))
        self.use_fisher: bool = use_fisher

        target_dev = original_weight.device if original_weight is not None else None
        target_dtype = original_weight.dtype if original_weight is not None else torch.float32

        if original_weight is not None:
            W = original_weight.float()
        else:
            W = torch.randn(out_features, in_features) * (1.0 / math.sqrt(in_features))

        # Fisher-weighted SVD
        if use_fisher and fisher_diag is not None and fisher_diag.shape[0] == in_features:
            scale = torch.sqrt(fisher_diag.float().clamp(min=1e-6) + 1e-6)
            # Normalize scale to avoid extreme values
            scale = scale / (scale.norm() / math.sqrt(in_features) + 1e-8)
            W_white = W * scale.unsqueeze(0)
        else:
            scale = None
            W_white = W

        U, S, Vt = torch.linalg.svd(W_white, full_matrices=False)

        # Unwhiten Vt
        if scale is not None:
            Vt_unw = Vt / scale.unsqueeze(0)
        else:
            Vt_unw = Vt

        r = self.rank

        self.register_buffer("U", U[:, :r].contiguous().to(device=target_dev, dtype=target_dtype))
        self.register_buffer(
            "Vt", Vt_unw[:r, :].contiguous().to(device=target_dev, dtype=target_dtype)
        )
        self.register_buffer("S", S[:r].contiguous().to(device=target_dev, dtype=target_dtype))

        self.register_buffer(
            "U_init", U[:, :r].contiguous().to(device=target_dev, dtype=target_dtype)
        )
        self.register_buffer(
            "Vt_init",
            Vt_unw[:r, :].contiguous().to(device=target_dev, dtype=target_dtype),
        )
        self.register_buffer("S_init", S[:r].contiguous().to(device=target_dev, dtype=target_dtype))

        self.gamma = nn.Parameter(torch.ones(r, device=target_dev, dtype=target_dtype))
        self.A = nn.Parameter(torch.zeros(out_features, r, device=target_dev, dtype=target_dtype))
        self.B = nn.Parameter(torch.zeros(r, in_features, device=target_dev, dtype=target_dtype))

        if original_bias is not None:
            self.bias = nn.Parameter(
                original_bias.clone().to(device=target_dev, dtype=target_dtype)
            )
        else:
            self.bias = nn.Parameter(
                torch.zeros(out_features, device=target_dev, dtype=target_dtype)
            )

    def init_gamma_from_sensitivity(self, sensitivity: torch.Tensor | None = None) -> None:
        """
        Initialize gamma from attention sensitivity per SVD direction.
        If sensitivity is None, initialize to ones.
        """
        if sensitivity is not None and sensitivity.shape[0] == self.rank:
            s_norm = sensitivity / (sensitivity.mean() + 1e-8)
            self.gamma.data.copy_(
                torch.sqrt(s_norm.clamp(0.1, 10.0)).to(
                    dtype=self.gamma.dtype, device=self.gamma.device
                )
            )
        else:
            self.gamma.data.fill_(1.0)

    def merge_to_linear(self) -> nn.Linear:
        """
        Merge SVD + gamma + LoRA into a single nn.Linear.
        After this, the layer has no LoRA overhead at inference.
        """
        with torch.no_grad():
            adapted_S = (self.S * self.gamma).to(self.U.dtype)
            W = (self.U * adapted_S.unsqueeze(0)) @ self.Vt
            W = W + self.A.to(self.U.dtype) @ self.B.to(self.U.dtype)
            bias = self.bias.clone() if self.bias is not None else None

        new_linear = nn.Linear(
            self.in_features,
            self.out_features,
            bias=bias is not None,
            device=W.device,
            dtype=W.dtype,
        )
        new_linear.weight.data.copy_(W.to(new_linear.weight.dtype))
        if bias is not None:
            new_linear.bias.data.copy_(bias.to(new_linear.bias.dtype))
        return new_linear

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective weight matrix (out_features, in_features)."""
        adapted_S = (self.S * self.gamma).to(self.U.dtype)
        W_svd = (self.U * adapted_S.unsqueeze(0)) @ self.Vt
        W_lora = self.A.to(self.U.dtype) @ self.B.to(self.U.dtype)
        return W_svd + W_lora

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        r"""Compute adapted linear projection efficiently in low-rank factorized form."""
        adapted_S = (self.S * self.gamma).to(dtype=x.dtype)
        # SVD branch: (x @ Vt.T) * adapted_S @ U.T
        x_vt = F.linear(x, self.Vt.to(dtype=x.dtype))
        x_svd = F.linear(x_vt * adapted_S, self.U.to(dtype=x.dtype))
        # LoRA branch: (x @ B.T) @ A.T
        x_lora = F.linear(F.linear(x, self.B.to(dtype=x.dtype)), self.A.to(dtype=x.dtype))
        out = x_svd + x_lora
        if self.bias is not None:
            out = out + self.bias.to(dtype=x.dtype)
        return out

    def freeze_base_weights(self) -> None:
        """Buffers U, S, Vt are non-trainable by construction. Gamma and A, B remain trainable."""
        pass

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

Linear projection layer with Fisher-weighted SVD initialization, attention-aware singular value modulation (\gamma), and low-rank residual branch (A B).

Formulation -----= $$W = U_{:, :r} \operatorname{diag}(\Sigma_{1:r} \odot \gamma) V_{1:r, :}^\top + A B$$

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective weight matrix (out_features, in_features)."""
    adapted_S = (self.S * self.gamma).to(self.U.dtype)
    W_svd = (self.U * adapted_S.unsqueeze(0)) @ self.Vt
    W_lora = self.A.to(self.U.dtype) @ self.B.to(self.U.dtype)
    return W_svd + W_lora

Reconstructed effective weight matrix (out_features, in_features).

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    r"""Compute adapted linear projection efficiently in low-rank factorized form."""
    adapted_S = (self.S * self.gamma).to(dtype=x.dtype)
    # SVD branch: (x @ Vt.T) * adapted_S @ U.T
    x_vt = F.linear(x, self.Vt.to(dtype=x.dtype))
    x_svd = F.linear(x_vt * adapted_S, self.U.to(dtype=x.dtype))
    # LoRA branch: (x @ B.T) @ A.T
    x_lora = F.linear(F.linear(x, self.B.to(dtype=x.dtype)), self.A.to(dtype=x.dtype))
    out = x_svd + x_lora
    if self.bias is not None:
        out = out + self.bias.to(dtype=x.dtype)
    return out

Compute adapted linear projection efficiently in low-rank factorized form.

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    """Buffers U, S, Vt are non-trainable by construction. Gamma and A, B remain trainable."""
    pass

Buffers U, S, Vt are non-trainable by construction. Gamma and A, B remain trainable.

def init_gamma_from_sensitivity(self, sensitivity: torch.Tensor | None = None) ‑> None

Expand source code

def init_gamma_from_sensitivity(self, sensitivity: torch.Tensor | None = None) -> None:
    """
    Initialize gamma from attention sensitivity per SVD direction.
    If sensitivity is None, initialize to ones.
    """
    if sensitivity is not None and sensitivity.shape[0] == self.rank:
        s_norm = sensitivity / (sensitivity.mean() + 1e-8)
        self.gamma.data.copy_(
            torch.sqrt(s_norm.clamp(0.1, 10.0)).to(
                dtype=self.gamma.dtype, device=self.gamma.device
            )
        )
    else:
        self.gamma.data.fill_(1.0)

Initialize gamma from attention sensitivity per SVD direction. If sensitivity is None, initialize to ones.

def merge_to_linear(self) ‑> torch.nn.modules.linear.Linear

Expand source code

def merge_to_linear(self) -> nn.Linear:
    """
    Merge SVD + gamma + LoRA into a single nn.Linear.
    After this, the layer has no LoRA overhead at inference.
    """
    with torch.no_grad():
        adapted_S = (self.S * self.gamma).to(self.U.dtype)
        W = (self.U * adapted_S.unsqueeze(0)) @ self.Vt
        W = W + self.A.to(self.U.dtype) @ self.B.to(self.U.dtype)
        bias = self.bias.clone() if self.bias is not None else None

    new_linear = nn.Linear(
        self.in_features,
        self.out_features,
        bias=bias is not None,
        device=W.device,
        dtype=W.dtype,
    )
    new_linear.weight.data.copy_(W.to(new_linear.weight.dtype))
    if bias is not None:
        new_linear.bias.data.copy_(bias.to(new_linear.bias.dtype))
    return new_linear

Merge SVD + gamma + LoRA into a single nn.Linear. After this, the layer has no LoRA overhead at inference.

class SAMBAM3Block (config: Any | None = None,
d_model: int | None = None,
n_heads: int | None = None,
window: int = 128,
rank: int | None = None,
pretrained_attn: nn.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("samba_m3", supports_causal=True, supports_gqa=True)
class SAMBAM3Block(nn.Module):
    r"""SAMBA-M3: Local Sliding-Window Attention + Spectral Mamba."""

    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        d_model: int | None = None,
        n_heads: int | None = None,
        window: int = 128,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__()
        if isinstance(config, int):
            d_model = config
            config = None

        # ---- d_model / n_heads ----
        if d_model is None:
            if config is not None:
                d_model = _cfg_get(config, "hidden_size", _cfg_get(config, "d_model", 768))
            elif pretrained_attn is not None:
                for attr in ("q_proj", "q_lin", "query"):
                    m = getattr(pretrained_attn, attr, None)
                    if m is not None and hasattr(m, "in_features"):
                        d_model = m.in_features
                        break
                if d_model is None:
                    d_model = getattr(pretrained_attn, "hidden_size", 768)
            else:
                d_model = kwargs.get("hidden_size", 768)

        if n_heads is None:
            n_heads = _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "n_heads", kwargs.get("num_attention_heads", 12)),
            )

        window = int(_cfg_get(config, "window_size", window))

        self.d_model: int = int(d_model)
        self.n_heads: int = int(n_heads)

        if self.n_heads <= 0 or self.d_model % self.n_heads != 0:
            raise ValueError(
                f"d_model ({self.d_model}) must be divisible by n_heads ({self.n_heads})"
            )

        self.head_dim: int = self.d_model // self.n_heads
        self.window: int = window
        self.lora_rank: int = int(kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank)))
        self.lora_alpha: float = float(
            kwargs.get("lora_alpha", _cfg_get(config, "lora_alpha", lora_alpha))
        )

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        self.num_kv_heads: int = int(cfg_kv) if cfg_kv is not None else self.n_heads
        self.num_key_value_groups: int = max(1, self.n_heads // max(1, self.num_kv_heads))

        self.q_head_size: int = self.n_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        # ---- SM3-FIX-08: decoder detection (override wins) ----
        if is_decoder is not None:
            is_dec = bool(is_decoder)
        else:
            model_type = str(_cfg_get(config, "model_type", "")).lower()
            is_dec = bool(
                _cfg_get(config, "is_decoder", False)
                or _cfg_get(config, "is_causal", False)
                or kwargs.get("is_causal", False)
            )
            decoder_prefixes = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            if any(model_type.startswith(p) for p in decoder_prefixes):
                is_dec = True
            arch = getattr(config, "architectures", None) if config is not None else None
            if arch and any("GPT" in str(a) or "CausalLM" in str(a) for a in arch):
                is_dec = True
        self.is_decoder: bool = bool(is_dec)
        self.is_causal: bool = self.is_decoder
        self.is_nested_encoder: bool = False

        # ---- Projections ----
        self.q_lin = LoRALinear(
            self.d_model,
            self.q_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.k_lin = LoRALinear(
            self.d_model,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.v_lin = LoRALinear(
            self.d_model,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.out_lin = LoRALinear(
            self.q_head_size,
            self.d_model,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        # ---- Mamba branch ----
        self.mamba = SpectralMamba(
            self.d_model, config=config, is_decoder=self.is_decoder, **kwargs
        )

        # ---- SM3-FIX-04: alpha always fp32, cast in forward ----
        self.alpha = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))
        self.gate_fn = nn.Sigmoid()

        dropout_p = float(
            _cfg_get(config, "dropout", _cfg_get(config, "attention_probs_dropout_prob", 0.1))
        )
        self.dropout = nn.Dropout(dropout_p)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    # ---- aliases ----
    @property
    def query(self):
        return self.q_lin

    @property
    def key(self):
        return self.k_lin

    @property
    def value(self):
        return self.v_lin

    @property
    def q_proj(self):
        return self.q_lin

    @property
    def k_proj(self):
        return self.k_lin

    @property
    def v_proj(self):
        return self.v_lin

    def get_ortho_loss(self) -> float:
        return 0.0

    def freeze_base_weights(self) -> None:
        for m in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
            m.freeze_base_weights()

    def unfreeze_base_weights(self) -> None:
        for m in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
            m.unfreeze_base_weights()

    # ------------------------------------------------------------------
    # Weight transfer (SM3-FIX-09)
    # ------------------------------------------------------------------
    def copy_weights(self, pretrained_attn: nn.Module) -> None:
        with torch.no_grad():
            w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

            src_cls_name = getattr(pretrained_attn, "__class__", type(pretrained_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(pretrained_attn, "query")
                and not hasattr(pretrained_attn, "dense")
                and not hasattr(pretrained_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(pretrained_attn, "self") and hasattr(pretrained_attn, "output"):
                self.is_nested_encoder = True
                src = pretrained_attn.self
                if hasattr(pretrained_attn.output, "dense"):
                    w_o = pretrained_attn.output.dense.weight
                    b_o = pretrained_attn.output.dense.bias
            else:
                src = pretrained_attn

            # --- GPT-2 fused ---
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                p_cls = getattr(src.c_attn, "__class__", type(src.c_attn)).__name__
                if (
                    "Conv1D" in p_cls
                    or hasattr(src.c_attn, "nf")
                    or (c_attn_w.shape[0] == self.d_model and c_attn_w.shape[1] != self.d_model)
                ):
                    c_attn_w = c_attn_w.t().contiguous()
                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size
                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]
                out_w = src.c_proj.weight.detach()
                p_proj_cls = getattr(src.c_proj, "__class__", type(src.c_proj)).__name__
                if (
                    "Conv1D" in p_proj_cls
                    or hasattr(src.c_proj, "nf")
                    or (out_w.shape[0] == self.d_model and out_w.shape[1] != self.d_model)
                ):
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # --- BERT/RoBERTa nested ---
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # --- Qwen / LLaMA ---
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = (
                    src.q_proj.bias.detach()
                    if getattr(src.q_proj, "bias", None) is not None
                    else None
                )
                w_k = src.k_proj.weight.detach()
                b_k = (
                    src.k_proj.bias.detach()
                    if getattr(src.k_proj, "bias", None) is not None
                    else None
                )
                w_v = src.v_proj.weight.detach()
                b_v = (
                    src.v_proj.bias.detach()
                    if getattr(src.v_proj, "bias", None) is not None
                    else None
                )
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = (
                        src.o_proj.bias.detach()
                        if getattr(src.o_proj, "bias", None) is not None
                        else None
                    )

            # --- Flat (DistilBERT) ---
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # --- Fallback ---
            if w_q is None:
                try:
                    adapter = UniversalAdapter()
                    for proj, target in (
                        ("q", self.q_lin),
                        ("k", self.k_lin),
                        ("v", self.v_lin),
                        ("o", self.out_lin),
                    ):
                        w, b = adapter.get_weight_bias(src, proj)
                        if w is None:
                            continue
                        if proj == "q":
                            w_q, b_q = w, b
                        elif proj == "k":
                            w_k, b_k = w, b
                        elif proj == "v":
                            w_v, b_v = w, b
                        else:
                            w_o, b_o = w, b
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # --- Copy into targets (SM3-FIX-09: cast + freeze) ---
            for target_mod, w_src, b_src, name in (
                (self.q_lin, w_q, b_q, "q_lin"),
                (self.k_lin, w_k, b_k, "k_lin"),
                (self.v_lin, w_v, b_v, "v_lin"),
                (self.out_lin, w_o, b_o, "out_lin"),
            ):
                if w_src is None:
                    continue
                tw = target_mod.weight
                w_cast = w_src.to(dtype=tw.dtype, device=tw.device)
                # orient if transposed
                if w_cast.shape != tw.shape and w_cast.t().contiguous().shape == tw.shape:
                    w_cast = w_cast.t().contiguous()
                min_r = min(tw.shape[0], w_cast.shape[0])
                min_c = min(tw.shape[1], w_cast.shape[1])
                tw.data[:min_r, :min_c].copy_(w_cast[:min_r, :min_c])
                if b_src is not None and target_mod.bias is not None:
                    tb = target_mod.bias
                    b_cast = b_src.to(dtype=tb.dtype, device=tb.device)
                    mb = min(tb.numel(), b_cast.numel())
                    tb.data[:mb].copy_(b_cast[:mb])
                fnorm = float(torch.linalg.norm(tw.data.float()))
                if fnorm < 1e-4:
                    logger.warning("SAMBAM3Block: %s norm too small (%.2e)", name, fnorm)
                target_mod.freeze_base_weights()

    def _chunked_sliding_window_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        past_len: int = 0,
    ) -> torch.Tensor:
        B, H, L, d = q.shape
        scale = 1.0 / math.sqrt(d)
        chunk = 64
        outputs = []

        for start in range(0, L, chunk):
            end = min(start + chunk, L)
            q_c = q[:, :, start:end]
            C_q = end - start

            # Absolute query positions (offset by past_len)
            q_abs_start = past_len + start
            q_abs_end = past_len + end

            # Key range: absolute positions in [max(0, q_abs_start - window), q_abs_end + window) for bidir
            if self.is_decoder:
                k_abs_start = max(0, q_abs_start - self.window)
                k_abs_end = q_abs_end
            else:
                k_abs_start = max(0, q_abs_start - self.window)
                k_abs_end = min(past_len + L, q_abs_end + self.window)

            # Slice keys in the *cache coordinates* (0..past_len+L)
            k_c = k[:, :, k_abs_start:k_abs_end]
            v_c = v[:, :, k_abs_start:k_abs_end]

            scores = torch.matmul(q_c.float(), k_c.transpose(-2, -1).contiguous().float()) * scale

            q_idx = torch.arange(q_abs_start, q_abs_end, device=q.device).view(1, 1, C_q, 1)
            k_idx = torch.arange(k_abs_start, k_abs_end, device=q.device).view(1, 1, 1, -1)
            dist = (q_idx - k_idx).abs()
            valid = dist <= self.window
            if self.is_decoder:
                valid = valid & (k_idx <= q_idx)

            scores = scores.masked_fill(~valid, -1e4)

            # SM3-FIX-07: head_mask dim==3 handled uniformly
            if attention_mask is not None:
                if attention_mask.dim() == 2:
                    # (B, N_total)
                    pad = (attention_mask[:, k_abs_start:k_abs_end] == 0).unsqueeze(1).unsqueeze(2)
                    scores = scores.masked_fill(pad, -1e4)
                elif attention_mask.dim() == 3:
                    scores = scores + attention_mask[:, start:end, k_abs_start:k_abs_end].unsqueeze(
                        1
                    ).to(scores.dtype)
                elif attention_mask.dim() == 4:
                    scores = scores + attention_mask[:, :, start:end, k_abs_start:k_abs_end].to(
                        scores.dtype
                    )

            probs = F.softmax(scores, dim=-1)
            probs = torch.nan_to_num(probs, nan=0.0)
            probs = self.dropout(probs.to(v_c.dtype))

            if head_mask is not None:
                hm = head_mask
                if hm.dim() == 1:
                    hm = hm.view(1, -1, 1, 1)
                elif hm.dim() == 2:
                    hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
                elif hm.dim() == 3:
                    # (B, H, Lk) — align to current key slice and unsqueeze
                    hm = hm[:, :, k_abs_start:k_abs_end].unsqueeze(2)
                hm = hm.to(probs.dtype)
                probs = probs * hm

            out_c = torch.matmul(probs, v_c)
            outputs.append(out_c)

        return torch.cat(outputs, dim=2)

    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,
        *args: Any,
        **kwargs: Any,
    ):
        if encoder_hidden_states is not None:
            raise NotImplementedError("SAMBAM3Block does not support cross-attention")

        B, L, D = hidden_states.shape
        target_dtype = hidden_states.dtype

        # ---- SM3-FIX-05: read cache ----
        past_k = past_v = None
        past_len = 0
        if past_key_value is not None:
            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]
                past_len = past_k.shape[-2]
            elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                past_k = past_key_value.key_cache[0]
                past_v = past_key_value.value_cache[0]
                past_len = past_k.shape[-2]

        # ---- QKV ----
        q = (
            self.q_lin(hidden_states)
            .view(B, L, self.n_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_lin(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_lin(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- SM3-FIX-06: RoPE with offset ----
        pos_emb = position_embeddings or kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
            cos, sin = pos_emb
            if cos.shape[-2] >= past_len + L:
                cos = cos.narrow(-2, past_len, L)
                sin = sin.narrow(-2, past_len, L)
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # ---- Append to KV cache ----
        if past_k is not None:
            k = torch.cat([past_k.to(k.dtype), k], dim=-2)
            v = torch.cat([past_v.to(v.dtype), v], dim=-2)

        present_key_value = (k, v)

        # ---- GQA expand ----
        if self.num_kv_heads != self.n_heads:
            k_exp = repeat_kv(k, self.num_key_value_groups)
            v_exp = repeat_kv(v, self.num_key_value_groups)
        else:
            k_exp, v_exp = k, v

        # ---- Attention ----
        attn_out = self._chunked_sliding_window_attention(
            q,
            k_exp,
            v_exp,
            attention_mask=attention_mask,
            head_mask=head_mask,
            past_len=past_len,
        )
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, L, self.q_head_size)
        attn_out = attn_out.to(target_dtype)

        # ---- Mamba branch ----
        mamba_out, _ = self.mamba(hidden_states, attention_mask=attention_mask, **kwargs)
        mamba_out = mamba_out.to(target_dtype)

        # ---- Fusion (SM3-FIX-04: alpha cast) ----
        alpha = torch.sigmoid(self.alpha).to(target_dtype)

        if self.is_nested_encoder:
            out = alpha * attn_out + (1.0 - alpha) * mamba_out
        else:
            attn_proj = self.out_lin(attn_out).to(target_dtype)
            out = alpha * attn_proj + (1.0 - alpha) * mamba_out

        out = torch.nan_to_num(out.to(target_dtype), nan=0.0)

        # ---- SM3-FIX-11: no silent pad; assert shape ----
        if out.shape[-1] != D:
            raise RuntimeError(f"SAMBAM3Block: output dim {out.shape[-1]} != expected {D}")

        return out, (present_key_value if output_attentions is False else (out, present_key_value))

SAMBA-M3: Local Sliding-Window Attention + Spectral Mamba.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

var supports_gqa : bool

The type of the None singleton.

Instance variables

prop k_proj

Expand source code

@property
def k_proj(self):
    return self.k_lin

prop key

Expand source code

@property
def key(self):
    return self.k_lin

prop q_proj

Expand source code

@property
def q_proj(self):
    return self.q_lin

prop query

Expand source code

@property
def query(self):
    return self.q_lin

prop v_proj

Expand source code

@property
def v_proj(self):
    return self.v_lin

prop value

Expand source code

@property
def value(self):
    return self.v_lin

Methods

def copy_weights(self, pretrained_attn: nn.Module) ‑> None

Expand source code

def copy_weights(self, pretrained_attn: nn.Module) -> None:
    with torch.no_grad():
        w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

        src_cls_name = getattr(pretrained_attn, "__class__", type(pretrained_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(pretrained_attn, "query")
            and not hasattr(pretrained_attn, "dense")
            and not hasattr(pretrained_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(pretrained_attn, "self") and hasattr(pretrained_attn, "output"):
            self.is_nested_encoder = True
            src = pretrained_attn.self
            if hasattr(pretrained_attn.output, "dense"):
                w_o = pretrained_attn.output.dense.weight
                b_o = pretrained_attn.output.dense.bias
        else:
            src = pretrained_attn

        # --- GPT-2 fused ---
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
            p_cls = getattr(src.c_attn, "__class__", type(src.c_attn)).__name__
            if (
                "Conv1D" in p_cls
                or hasattr(src.c_attn, "nf")
                or (c_attn_w.shape[0] == self.d_model and c_attn_w.shape[1] != self.d_model)
            ):
                c_attn_w = c_attn_w.t().contiguous()
            q_dim = self.q_head_size
            k_dim = self.kv_head_size
            v_dim = self.kv_head_size
            w_q = c_attn_w[:q_dim, :]
            w_k = c_attn_w[q_dim : q_dim + k_dim, :]
            w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
            if c_attn_b is not None:
                b_q = c_attn_b[:q_dim]
                b_k = c_attn_b[q_dim : q_dim + k_dim]
                b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]
            out_w = src.c_proj.weight.detach()
            p_proj_cls = getattr(src.c_proj, "__class__", type(src.c_proj)).__name__
            if (
                "Conv1D" in p_proj_cls
                or hasattr(src.c_proj, "nf")
                or (out_w.shape[0] == self.d_model and out_w.shape[1] != self.d_model)
            ):
                out_w = out_w.t().contiguous()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # --- BERT/RoBERTa nested ---
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None
            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # --- Qwen / LLaMA ---
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = (
                src.q_proj.bias.detach()
                if getattr(src.q_proj, "bias", None) is not None
                else None
            )
            w_k = src.k_proj.weight.detach()
            b_k = (
                src.k_proj.bias.detach()
                if getattr(src.k_proj, "bias", None) is not None
                else None
            )
            w_v = src.v_proj.weight.detach()
            b_v = (
                src.v_proj.bias.detach()
                if getattr(src.v_proj, "bias", None) is not None
                else None
            )
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = (
                    src.o_proj.bias.detach()
                    if getattr(src.o_proj, "bias", None) is not None
                    else None
                )

        # --- Flat (DistilBERT) ---
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # --- Fallback ---
        if w_q is None:
            try:
                adapter = UniversalAdapter()
                for proj, target in (
                    ("q", self.q_lin),
                    ("k", self.k_lin),
                    ("v", self.v_lin),
                    ("o", self.out_lin),
                ):
                    w, b = adapter.get_weight_bias(src, proj)
                    if w is None:
                        continue
                    if proj == "q":
                        w_q, b_q = w, b
                    elif proj == "k":
                        w_k, b_k = w, b
                    elif proj == "v":
                        w_v, b_v = w, b
                    else:
                        w_o, b_o = w, b
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # --- Copy into targets (SM3-FIX-09: cast + freeze) ---
        for target_mod, w_src, b_src, name in (
            (self.q_lin, w_q, b_q, "q_lin"),
            (self.k_lin, w_k, b_k, "k_lin"),
            (self.v_lin, w_v, b_v, "v_lin"),
            (self.out_lin, w_o, b_o, "out_lin"),
        ):
            if w_src is None:
                continue
            tw = target_mod.weight
            w_cast = w_src.to(dtype=tw.dtype, device=tw.device)
            # orient if transposed
            if w_cast.shape != tw.shape and w_cast.t().contiguous().shape == tw.shape:
                w_cast = w_cast.t().contiguous()
            min_r = min(tw.shape[0], w_cast.shape[0])
            min_c = min(tw.shape[1], w_cast.shape[1])
            tw.data[:min_r, :min_c].copy_(w_cast[:min_r, :min_c])
            if b_src is not None and target_mod.bias is not None:
                tb = target_mod.bias
                b_cast = b_src.to(dtype=tb.dtype, device=tb.device)
                mb = min(tb.numel(), b_cast.numel())
                tb.data[:mb].copy_(b_cast[:mb])
            fnorm = float(torch.linalg.norm(tw.data.float()))
            if fnorm < 1e-4:
                logger.warning("SAMBAM3Block: %s norm too small (%.2e)", name, fnorm)
            target_mod.freeze_base_weights()

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,
*args: Any,
**kwargs: Any) ‑> Callable[..., typing.Any]

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,
    *args: Any,
    **kwargs: Any,
):
    if encoder_hidden_states is not None:
        raise NotImplementedError("SAMBAM3Block does not support cross-attention")

    B, L, D = hidden_states.shape
    target_dtype = hidden_states.dtype

    # ---- SM3-FIX-05: read cache ----
    past_k = past_v = None
    past_len = 0
    if past_key_value is not None:
        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]
            past_len = past_k.shape[-2]
        elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
            past_k = past_key_value.key_cache[0]
            past_v = past_key_value.value_cache[0]
            past_len = past_k.shape[-2]

    # ---- QKV ----
    q = (
        self.q_lin(hidden_states)
        .view(B, L, self.n_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    k = (
        self.k_lin(hidden_states)
        .view(B, L, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v = (
        self.v_lin(hidden_states)
        .view(B, L, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # ---- SM3-FIX-06: RoPE with offset ----
    pos_emb = position_embeddings or kwargs.get("position_embeddings")
    if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
        cos, sin = pos_emb
        if cos.shape[-2] >= past_len + L:
            cos = cos.narrow(-2, past_len, L)
            sin = sin.narrow(-2, past_len, L)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    # ---- Append to KV cache ----
    if past_k is not None:
        k = torch.cat([past_k.to(k.dtype), k], dim=-2)
        v = torch.cat([past_v.to(v.dtype), v], dim=-2)

    present_key_value = (k, v)

    # ---- GQA expand ----
    if self.num_kv_heads != self.n_heads:
        k_exp = repeat_kv(k, self.num_key_value_groups)
        v_exp = repeat_kv(v, self.num_key_value_groups)
    else:
        k_exp, v_exp = k, v

    # ---- Attention ----
    attn_out = self._chunked_sliding_window_attention(
        q,
        k_exp,
        v_exp,
        attention_mask=attention_mask,
        head_mask=head_mask,
        past_len=past_len,
    )
    attn_out = attn_out.transpose(1, 2).contiguous().view(B, L, self.q_head_size)
    attn_out = attn_out.to(target_dtype)

    # ---- Mamba branch ----
    mamba_out, _ = self.mamba(hidden_states, attention_mask=attention_mask, **kwargs)
    mamba_out = mamba_out.to(target_dtype)

    # ---- Fusion (SM3-FIX-04: alpha cast) ----
    alpha = torch.sigmoid(self.alpha).to(target_dtype)

    if self.is_nested_encoder:
        out = alpha * attn_out + (1.0 - alpha) * mamba_out
    else:
        attn_proj = self.out_lin(attn_out).to(target_dtype)
        out = alpha * attn_proj + (1.0 - alpha) * mamba_out

    out = torch.nan_to_num(out.to(target_dtype), nan=0.0)

    # ---- SM3-FIX-11: no silent pad; assert shape ----
    if out.shape[-1] != D:
        raise RuntimeError(f"SAMBAM3Block: output dim {out.shape[-1]} != expected {D}")

    return out, (present_key_value if output_attentions is False else (out, present_key_value))

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    for m in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
        m.freeze_base_weights()

def get_ortho_loss(self) ‑> float

Expand source code

def get_ortho_loss(self) -> float:
    return 0.0

def unfreeze_base_weights(self) ‑> None

Expand source code

def unfreeze_base_weights(self) -> None:
    for m in (self.q_lin, self.k_lin, self.v_lin, self.out_lin):
        m.unfreeze_base_weights()

class SOAAttention (config: Any | None = None,
model_config: Any | None = None,
rank: int | None = None,
k_rank: int = 64,
dropout: float = 0.1,
pretrained_attn: Any | None = None,
old_attn: Any | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("soa_module", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("soa_attn")
@ATTENTION_REGISTRY.register("soa_fused_module")
@ATTENTION_REGISTRY.register("soafusedattention")
@ATTENTION_REGISTRY.register("soa_fused_attn")
class SOAAttention(BaseAttention):
    r"""Shared Orthogonal Attention (SOA) with Per-Projection Latent Bottlenecks.

    Factorizes dense Q, K, and V linear projections into low-rank bottleneck pairs
    initialized via truncated SVD from pretrained transformer weights.
    """

    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        rank: int | None = None,
        k_rank: int = 64,
        dropout: float = 0.1,
        pretrained_attn: Any | None = None,
        old_attn: Any | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config or model_config or kwargs.get("config")
        src_attn = (
            pretrained_attn
            if pretrained_attn is not None
            else (old_attn if old_attn is not None else kwargs.get("source_attn"))
        )
        super().__init__(cfg, rank=k_rank, **kwargs)

        # 1. Architecture detection & dimensions
        self.hidden_size: int = (
            getattr(
                cfg,
                "hidden_size",
                getattr(cfg, "dim", kwargs.get("hidden_size", kwargs.get("dim", 768))),
            )
            if cfg is not None
            else kwargs.get("hidden_size", kwargs.get("dim", 768))
        )

        self.num_heads: int = (
            getattr(
                cfg,
                "num_attention_heads",
                getattr(
                    cfg,
                    "n_heads",
                    getattr(
                        cfg,
                        "num_heads",
                        kwargs.get("num_attention_heads", kwargs.get("num_heads", 12)),
                    ),
                ),
            )
            if cfg is not None
            else kwargs.get("num_attention_heads", kwargs.get("num_heads", 12))
        )
        self.num_heads = int(self.num_heads or 12)

        # Probe head_dim from source attention if available
        head_dim = kwargs.get("head_dim")
        if not head_dim and src_attn is not None:
            for attr in ("q_proj", "query", "q_lin"):
                mod = getattr(src_attn, attr, None)
                if mod is not None and hasattr(mod, "weight") and mod.weight is not None:
                    if mod.weight.ndim == 2:
                        head_dim = mod.weight.shape[0] // max(1, self.num_heads)
                        break
            if not head_dim:
                head_dim = getattr(src_attn, "head_dim", None)
        if not head_dim and cfg is not None:
            head_dim = getattr(cfg, "head_dim", None)
        if not head_dim or head_dim <= 0:
            head_dim = self.hidden_size // max(1, self.num_heads)
        self.head_dim = int(head_dim)

        cfg_kv = (
            getattr(
                cfg,
                "num_key_value_heads",
                getattr(
                    cfg,
                    "num_kv_heads",
                    kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads")),
                ),
            )
            if cfg is not None
            else kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads"))
        )
        self.num_kv_heads: int = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_key_value_groups: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.kv_repeat: int = self.num_key_value_groups

        self.q_head_size: int = self.num_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        # Adaptive k_rank
        comp_config = getattr(cfg, "compression_config", None)
        if k_rank is None or k_rank == 64:
            effective_rank = getattr(
                comp_config,
                "k_rank",
                getattr(
                    comp_config,
                    "soa_rank",
                    rank if rank is not None else max(128, self.hidden_size // 4),
                ),
            )
        else:
            effective_rank = min(k_rank, self.hidden_size)
        self.k_rank: int = effective_rank
        self.rank = effective_rank
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha

        # Decoder / Causal flag
        _mt = str(getattr(cfg, "model_type", "") or "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_dec = (
            bool(_cfg_get(cfg, "is_decoder", False))
            or bool(_cfg_get(cfg, "is_causal", False))
            or any(_mt.startswith(p) for p in _decoder_prefixes)
            or (src_attn is not None and hasattr(src_attn, "c_attn"))
            or (src_attn is not None and getattr(src_attn, "is_causal", False))
            or (src_attn is not None and getattr(src_attn, "is_decoder", False))
            or bool(kwargs.get("is_causal", False))
            or bool(kwargs.get("is_decoder", False))
        )
        self.is_decoder: bool = is_dec
        self.is_causal: bool = is_dec
        self.is_nested_encoder: bool = False

        # 2. Factorized SVD linear projections for Q, K, V
        self.q_down = nn.Linear(self.hidden_size, self.k_rank, bias=False)
        self.q_up = FactorizedLoRAUp(
            self.k_rank,
            self.q_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        self.k_down = nn.Linear(self.hidden_size, self.k_rank, bias=False)
        self.k_up = FactorizedLoRAUp(
            self.k_rank,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        self.v_down = nn.Linear(self.hidden_size, self.k_rank, bias=False)
        self.v_up = FactorizedLoRAUp(
            self.k_rank,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        # Output projection with trainable LoRA adapter
        self.out_proj = LoRALinear(
            self.q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", dropout))
            if cfg is not None
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and src_attn is not None:
            target_dtype = next(src_attn.parameters(), torch.tensor([])).dtype
            if not isinstance(target_dtype, torch.dtype):
                target_dtype = torch.float32

        # 3. Pretrained weight transfer
        if src_attn is not None:
            self.copy_weights(src_attn)

        if target_dtype is not None:
            self.to(target_dtype)

    @property
    def q_proj(self) -> FactorizedLinear:
        return FactorizedLinear(self.q_down, self.q_up)

    @property
    def k_proj(self) -> FactorizedLinear:
        return FactorizedLinear(self.k_down, self.k_up)

    @property
    def v_proj(self) -> FactorizedLinear:
        return FactorizedLinear(self.v_down, self.v_up)

    @property
    def query(self) -> FactorizedLinear:
        return self.q_proj

    @property
    def key(self) -> FactorizedLinear:
        return self.k_proj

    @property
    def value(self) -> FactorizedLinear:
        return self.v_proj

    def copy_weights(self, old_attn: nn.Module) -> None:
        """Robust weight extraction and SVD factorization across all architectures."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Check nested encoder container
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_out_lin = (
                hasattr(old_attn, "out_lin")
                or hasattr(old_attn, "c_proj")
                or hasattr(old_attn, "o_proj")
                or hasattr(old_attn, "dense")
            )
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query") and not has_out_lin
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                self.is_decoder = True
                self.is_causal = True

                c_attn_mod = src.c_attn
                base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
                c_attn_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)).detach()
                c_attn_b = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                c_attn_b = c_attn_b.detach() if c_attn_b is not None else None

                c_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in c_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (c_attn_w.dim() == 2 and c_attn_w.shape[0] < c_attn_w.shape[1])
                ):
                    c_attn_w = c_attn_w.t().contiguous()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size

                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

                c_proj_mod = src.c_proj
                base_p = getattr(c_proj_mod, "base_layer", c_proj_mod)
                out_w = getattr(base_p, "weight", getattr(c_proj_mod, "weight", None)).detach()
                b_o = getattr(base_p, "bias", getattr(c_proj_mod, "bias", None))
                b_o = b_o.detach() if b_o is not None else None

                p_cls = getattr(base_p, "__class__", type(base_p)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_p, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (out_w.dim() == 2 and out_w.shape[0] < out_w.shape[1])
                ):
                    out_w = out_w.t().contiguous()
                w_o = out_w

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None

                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None

                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # Pattern 5: UniversalAdapter fallback
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    w_q, b_q = adapter.get_weight_bias(src, "q")
                    w_k, b_k = adapter.get_weight_bias(src, "k")
                    w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception:
                    pass

            # Strict validation
            missing = []
            if w_q is None:
                missing.append("Q projection")
            if w_k is None:
                missing.append("K projection")
            if w_v is None:
                missing.append("V projection")
            if w_o is None and not self.is_nested_encoder:
                missing.append("Out projection")

            if missing:
                raise RuntimeError(
                    f"SOAAttention: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
                )

            # SVD Factorization helper
            def _factorize_svd(
                W: torch.Tensor,
                down: nn.Linear,
                up: nn.Linear,
                bias: torch.Tensor | None,
            ) -> None:
                W_f = W.float()
                U, S, Vt = torch.linalg.svd(W_f, full_matrices=False)
                S = S.clamp(min=1e-6)
                r = min(self.k_rank, U.size(1), Vt.size(0))

                down.weight.zero_()
                up.weight.zero_()
                sqrt_s = torch.sqrt(S[:r])
                down.weight[:r, :].copy_((sqrt_s.unsqueeze(1) * Vt[:r, :]).to(down.weight.dtype))
                up.weight[:, :r].copy_((U[:, :r] * sqrt_s).to(up.weight.dtype))

                if bias is not None and up.bias is not None:
                    up.bias.copy_(bias.to(up.bias.dtype))
                elif up.bias is not None:
                    up.bias.zero_()

                # Frobenius check
                rec = up.weight @ down.weight
                f_norm = float(torch.linalg.norm(rec))
                if f_norm < 1e-4:
                    logger.warning(
                        f"SOAAttention: SVD reconstructed weight norm is near zero: {f_norm}"
                    )

            _factorize_svd(w_q, self.q_down, self.q_up, b_q)
            _factorize_svd(w_k, self.k_down, self.k_up, b_k)
            _factorize_svd(w_v, self.v_down, self.v_up, b_v)

            if w_o is not None and not self.is_nested_encoder:
                self.out_proj.weight.copy_(w_o.to(self.out_proj.weight.dtype))
                if b_o is not None and self.out_proj.bias is not None:
                    self.out_proj.bias.copy_(b_o.to(self.out_proj.bias.dtype))
                f_norm_o = float(torch.linalg.norm(self.out_proj.weight))
                if f_norm_o < 1e-4:
                    logger.warning(f"SOAAttention: Out projection norm is near zero: {f_norm_o}")
            elif self.is_nested_encoder:
                if hasattr(self.out_proj, "weight") and self.out_proj.weight.ndim == 2:
                    nn.init.eye_(self.out_proj.weight)
                if hasattr(self.out_proj, "bias") and self.out_proj.bias is not None:
                    nn.init.zeros_(self.out_proj.bias)

            # Freeze base weights
            for mod in [self.q_down, self.k_down, self.v_down]:
                mod.weight.requires_grad = False
                if getattr(mod, "bias", None) is not None:
                    mod.bias.requires_grad = False
            for mod in [self.q_up, self.k_up, self.v_up]:
                if hasattr(mod, "freeze_base_weights"):
                    mod.freeze_base_weights()
                else:
                    mod.weight.requires_grad = False
                    if getattr(mod, "bias", None) is not None:
                        mod.bias.requires_grad = False
                    if getattr(mod, "lora_A", None) is not None:
                        mod.lora_A.requires_grad = True
                    if getattr(mod, "lora_B", None) is not None:
                        mod.lora_B.requires_grad = True
            if hasattr(self.out_proj, "freeze_base_weights"):
                self.out_proj.freeze_base_weights()
            else:
                self.out_proj.weight.requires_grad = False
                if getattr(self.out_proj, "bias", None) is not None:
                    self.out_proj.bias.requires_grad = False
                if getattr(self.out_proj, "lora_A", None) is not None:
                    self.out_proj.lora_A.requires_grad = True
                if getattr(self.out_proj, "lora_B", None) is not None:
                    self.out_proj.lora_B.requires_grad = True

    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> "SOAAttention":
        return cls(model_config=model_config, pretrained_attn=old_attn, **kwargs)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        B, N, _ = hidden_states.shape

        # Factorized projections
        z_q = self.q_down(hidden_states)
        z_k = self.k_down(hidden_states)
        z_v = self.v_down(hidden_states)

        q = self.q_up(z_q)
        k = self.k_up(z_k)
        v = self.v_up(z_v)

        H = self.num_heads
        H_kv = self.num_kv_heads
        D = self.head_dim

        q = q.view(B, N, H, D).transpose(1, 2).contiguous()
        k = k.view(B, N, H_kv, D).transpose(1, 2).contiguous()
        v = v.view(B, N, H_kv, D).transpose(1, 2).contiguous()

        # Partial Rotary Position Embeddings (RoPE)
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
            cos, sin = pos_emb
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # Grouped-Query Attention expansion
        if H_kv != H:
            k = repeat_kv(k, self.num_key_value_groups)
            v = repeat_kv(v, self.num_key_value_groups)

        scale = 1.0 / math.sqrt(D)
        # Memory-efficient: use SDPA to avoid OOM on large seqs

        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

            _sdpa_out = torch.nn.functional.scaled_dot_product_attention(
                _qs,
                _ks,
                _vs,
                is_causal=getattr(self, "is_decoder", False) and N > 1,
                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, self.all_head_size)
            )

            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None

        except Exception:
            pass

        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        # Causal mask for decoders
        is_causal = (
            self.is_decoder or getattr(self, "is_causal", False) or kwargs.get("is_causal", False)
        ) and (encoder_hidden_states is None)
        if is_causal and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full((N, N), float("-inf"), device=scores.device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        # Standard Attention Mask
        mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask")
        if mask is not None:
            if mask.dim() == 4 and mask.dtype == torch.bool:
                scores = scores.masked_fill(~mask, -10000.0)
            else:
                scores = apply_attention_mask(scores, mask)

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights.to(v.dtype))

        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)
            attn_weights = attn_weights * head_mask

        context = torch.matmul(attn_weights, v)
        context = context.transpose(1, 2).contiguous().view(B, N, H * D)

        if self.is_nested_encoder:
            return (context, attn_weights if output_attentions else None)

        output = self.out_proj(context)
        return (output, attn_weights if output_attentions else None)

Shared Orthogonal Attention (SOA) with Per-Projection Latent Bottlenecks.

Factorizes dense Q, K, and V linear projections into low-rank bottleneck pairs initialized via truncated SVD from pretrained transformer weights.

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_gqa : bool

The type of the None singleton.

Static methods

def from_distilbert_pretrained(old_attn: torch.nn.modules.module.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> SOAAttention

Instance variables

prop k_proj : FactorizedLinear

Expand source code

@property
def k_proj(self) -> FactorizedLinear:
    return FactorizedLinear(self.k_down, self.k_up)

prop key : FactorizedLinear

Expand source code

@property
def key(self) -> FactorizedLinear:
    return self.k_proj

prop q_proj : FactorizedLinear

Expand source code

@property
def q_proj(self) -> FactorizedLinear:
    return FactorizedLinear(self.q_down, self.q_up)

prop query : FactorizedLinear

Expand source code

@property
def query(self) -> FactorizedLinear:
    return self.q_proj

prop v_proj : FactorizedLinear

Expand source code

@property
def v_proj(self) -> FactorizedLinear:
    return FactorizedLinear(self.v_down, self.v_up)

prop value : FactorizedLinear

Expand source code

@property
def value(self) -> FactorizedLinear:
    return self.v_proj

Methods

def copy_weights(self, old_attn: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    """Robust weight extraction and SVD factorization across all architectures."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Check nested encoder container
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        has_out_lin = (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
            or hasattr(old_attn, "dense")
        )
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query") and not has_out_lin
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight
                b_o = old_attn.output.dense.bias
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            self.is_decoder = True
            self.is_causal = True

            c_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            c_attn_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)).detach()
            c_attn_b = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            c_attn_b = c_attn_b.detach() if c_attn_b is not None else None

            c_cls = getattr(base_c, "__class__", type(base_c)).__name__
            if (
                ("Conv1D" in c_cls)
                or hasattr(base_c, "nf")
                or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                or hasattr(c_attn_mod, "nf")
                or (c_attn_w.dim() == 2 and c_attn_w.shape[0] < c_attn_w.shape[1])
            ):
                c_attn_w = c_attn_w.t().contiguous()

            q_dim = self.q_head_size
            k_dim = self.kv_head_size
            v_dim = self.kv_head_size

            w_q = c_attn_w[:q_dim, :]
            w_k = c_attn_w[q_dim : q_dim + k_dim, :]
            w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

            if c_attn_b is not None:
                b_q = c_attn_b[:q_dim]
                b_k = c_attn_b[q_dim : q_dim + k_dim]
                b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

            c_proj_mod = src.c_proj
            base_p = getattr(c_proj_mod, "base_layer", c_proj_mod)
            out_w = getattr(base_p, "weight", getattr(c_proj_mod, "weight", None)).detach()
            b_o = getattr(base_p, "bias", getattr(c_proj_mod, "bias", None))
            b_o = b_o.detach() if b_o is not None else None

            p_cls = getattr(base_p, "__class__", type(base_p)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(base_p, "nf")
                or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                or hasattr(c_proj_mod, "nf")
                or (out_w.dim() == 2 and out_w.shape[0] < out_w.shape[1])
            ):
                out_w = out_w.t().contiguous()
            w_o = out_w

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None

            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None

            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # Pattern 5: UniversalAdapter fallback
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                w_q, b_q = adapter.get_weight_bias(src, "q")
                w_k, b_k = adapter.get_weight_bias(src, "k")
                w_v, b_v = adapter.get_weight_bias(src, "v")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(src, "o")
            except Exception:
                pass

        # Strict validation
        missing = []
        if w_q is None:
            missing.append("Q projection")
        if w_k is None:
            missing.append("K projection")
        if w_v is None:
            missing.append("V projection")
        if w_o is None and not self.is_nested_encoder:
            missing.append("Out projection")

        if missing:
            raise RuntimeError(
                f"SOAAttention: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
            )

        # SVD Factorization helper
        def _factorize_svd(
            W: torch.Tensor,
            down: nn.Linear,
            up: nn.Linear,
            bias: torch.Tensor | None,
        ) -> None:
            W_f = W.float()
            U, S, Vt = torch.linalg.svd(W_f, full_matrices=False)
            S = S.clamp(min=1e-6)
            r = min(self.k_rank, U.size(1), Vt.size(0))

            down.weight.zero_()
            up.weight.zero_()
            sqrt_s = torch.sqrt(S[:r])
            down.weight[:r, :].copy_((sqrt_s.unsqueeze(1) * Vt[:r, :]).to(down.weight.dtype))
            up.weight[:, :r].copy_((U[:, :r] * sqrt_s).to(up.weight.dtype))

            if bias is not None and up.bias is not None:
                up.bias.copy_(bias.to(up.bias.dtype))
            elif up.bias is not None:
                up.bias.zero_()

            # Frobenius check
            rec = up.weight @ down.weight
            f_norm = float(torch.linalg.norm(rec))
            if f_norm < 1e-4:
                logger.warning(
                    f"SOAAttention: SVD reconstructed weight norm is near zero: {f_norm}"
                )

        _factorize_svd(w_q, self.q_down, self.q_up, b_q)
        _factorize_svd(w_k, self.k_down, self.k_up, b_k)
        _factorize_svd(w_v, self.v_down, self.v_up, b_v)

        if w_o is not None and not self.is_nested_encoder:
            self.out_proj.weight.copy_(w_o.to(self.out_proj.weight.dtype))
            if b_o is not None and self.out_proj.bias is not None:
                self.out_proj.bias.copy_(b_o.to(self.out_proj.bias.dtype))
            f_norm_o = float(torch.linalg.norm(self.out_proj.weight))
            if f_norm_o < 1e-4:
                logger.warning(f"SOAAttention: Out projection norm is near zero: {f_norm_o}")
        elif self.is_nested_encoder:
            if hasattr(self.out_proj, "weight") and self.out_proj.weight.ndim == 2:
                nn.init.eye_(self.out_proj.weight)
            if hasattr(self.out_proj, "bias") and self.out_proj.bias is not None:
                nn.init.zeros_(self.out_proj.bias)

        # Freeze base weights
        for mod in [self.q_down, self.k_down, self.v_down]:
            mod.weight.requires_grad = False
            if getattr(mod, "bias", None) is not None:
                mod.bias.requires_grad = False
        for mod in [self.q_up, self.k_up, self.v_up]:
            if hasattr(mod, "freeze_base_weights"):
                mod.freeze_base_weights()
            else:
                mod.weight.requires_grad = False
                if getattr(mod, "bias", None) is not None:
                    mod.bias.requires_grad = False
                if getattr(mod, "lora_A", None) is not None:
                    mod.lora_A.requires_grad = True
                if getattr(mod, "lora_B", None) is not None:
                    mod.lora_B.requires_grad = True
        if hasattr(self.out_proj, "freeze_base_weights"):
            self.out_proj.freeze_base_weights()
        else:
            self.out_proj.weight.requires_grad = False
            if getattr(self.out_proj, "bias", None) is not None:
                self.out_proj.bias.requires_grad = False
            if getattr(self.out_proj, "lora_A", None) is not None:
                self.out_proj.lora_A.requires_grad = True
            if getattr(self.out_proj, "lora_B", None) is not None:
                self.out_proj.lora_B.requires_grad = True

Robust weight extraction and SVD factorization across all architectures.

def forward(self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
head_mask: torch.Tensor | None = None,
encoder_hidden_states: torch.Tensor | None = None,
encoder_attention_mask: torch.Tensor | None = None,
past_key_value: Any | None = None,
output_attentions: bool | None = False,
output_hidden_states: bool | None = False,
return_dict: bool | None = True,
*args: Any,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor

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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
    B, N, _ = hidden_states.shape

    # Factorized projections
    z_q = self.q_down(hidden_states)
    z_k = self.k_down(hidden_states)
    z_v = self.v_down(hidden_states)

    q = self.q_up(z_q)
    k = self.k_up(z_k)
    v = self.v_up(z_v)

    H = self.num_heads
    H_kv = self.num_kv_heads
    D = self.head_dim

    q = q.view(B, N, H, D).transpose(1, 2).contiguous()
    k = k.view(B, N, H_kv, D).transpose(1, 2).contiguous()
    v = v.view(B, N, H_kv, D).transpose(1, 2).contiguous()

    # Partial Rotary Position Embeddings (RoPE)
    pos_emb = kwargs.get("position_embeddings")
    if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
        cos, sin = pos_emb
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    # Grouped-Query Attention expansion
    if H_kv != H:
        k = repeat_kv(k, self.num_key_value_groups)
        v = repeat_kv(v, self.num_key_value_groups)

    scale = 1.0 / math.sqrt(D)
    # Memory-efficient: use SDPA to avoid OOM on large seqs

    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

        _sdpa_out = torch.nn.functional.scaled_dot_product_attention(
            _qs,
            _ks,
            _vs,
            is_causal=getattr(self, "is_decoder", False) and N > 1,
            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, self.all_head_size)
        )

        _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

        return (
            self.out_proj(_sdpa_out)
            if not getattr(self, "is_nested_encoder", False)
            else _sdpa_out
        ), None

    except Exception:
        pass

    scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

    # Causal mask for decoders
    is_causal = (
        self.is_decoder or getattr(self, "is_causal", False) or kwargs.get("is_causal", False)
    ) and (encoder_hidden_states is None)
    if is_causal and N > 1:
        causal_mask = (
            torch.triu(
                torch.full((N, N), float("-inf"), device=scores.device, dtype=torch.float32),
                diagonal=1,
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )
        scores = scores + causal_mask

    # Standard Attention Mask
    mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask")
    if mask is not None:
        if mask.dim() == 4 and mask.dtype == torch.bool:
            scores = scores.masked_fill(~mask, -10000.0)
        else:
            scores = apply_attention_mask(scores, mask)

    attn_weights = F.softmax(scores, dim=-1)
    attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
    attn_weights = self.dropout(attn_weights.to(v.dtype))

    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)
        attn_weights = attn_weights * head_mask

    context = torch.matmul(attn_weights, v)
    context = context.transpose(1, 2).contiguous().view(B, N, H * D)

    if self.is_nested_encoder:
        return (context, attn_weights if output_attentions else None)

    output = self.out_proj(context)
    return (output, attn_weights if output_attentions else None)

Must return (output, attention_weights_or_None).

class SOAFusedAttention (config: Any | None = None,
model_config: Any | None = None,
rank: int | None = None,
k_rank: int = 64,
dropout: float = 0.1,
pretrained_attn: Any | None = None,
old_attn: Any | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("soa_module", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("soa_attn")
@ATTENTION_REGISTRY.register("soa_fused_module")
@ATTENTION_REGISTRY.register("soafusedattention")
@ATTENTION_REGISTRY.register("soa_fused_attn")
class SOAAttention(BaseAttention):
    r"""Shared Orthogonal Attention (SOA) with Per-Projection Latent Bottlenecks.

    Factorizes dense Q, K, and V linear projections into low-rank bottleneck pairs
    initialized via truncated SVD from pretrained transformer weights.
    """

    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        rank: int | None = None,
        k_rank: int = 64,
        dropout: float = 0.1,
        pretrained_attn: Any | None = None,
        old_attn: Any | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        cfg = config or model_config or kwargs.get("config")
        src_attn = (
            pretrained_attn
            if pretrained_attn is not None
            else (old_attn if old_attn is not None else kwargs.get("source_attn"))
        )
        super().__init__(cfg, rank=k_rank, **kwargs)

        # 1. Architecture detection & dimensions
        self.hidden_size: int = (
            getattr(
                cfg,
                "hidden_size",
                getattr(cfg, "dim", kwargs.get("hidden_size", kwargs.get("dim", 768))),
            )
            if cfg is not None
            else kwargs.get("hidden_size", kwargs.get("dim", 768))
        )

        self.num_heads: int = (
            getattr(
                cfg,
                "num_attention_heads",
                getattr(
                    cfg,
                    "n_heads",
                    getattr(
                        cfg,
                        "num_heads",
                        kwargs.get("num_attention_heads", kwargs.get("num_heads", 12)),
                    ),
                ),
            )
            if cfg is not None
            else kwargs.get("num_attention_heads", kwargs.get("num_heads", 12))
        )
        self.num_heads = int(self.num_heads or 12)

        # Probe head_dim from source attention if available
        head_dim = kwargs.get("head_dim")
        if not head_dim and src_attn is not None:
            for attr in ("q_proj", "query", "q_lin"):
                mod = getattr(src_attn, attr, None)
                if mod is not None and hasattr(mod, "weight") and mod.weight is not None:
                    if mod.weight.ndim == 2:
                        head_dim = mod.weight.shape[0] // max(1, self.num_heads)
                        break
            if not head_dim:
                head_dim = getattr(src_attn, "head_dim", None)
        if not head_dim and cfg is not None:
            head_dim = getattr(cfg, "head_dim", None)
        if not head_dim or head_dim <= 0:
            head_dim = self.hidden_size // max(1, self.num_heads)
        self.head_dim = int(head_dim)

        cfg_kv = (
            getattr(
                cfg,
                "num_key_value_heads",
                getattr(
                    cfg,
                    "num_kv_heads",
                    kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads")),
                ),
            )
            if cfg is not None
            else kwargs.get("num_key_value_heads", kwargs.get("num_kv_heads"))
        )
        self.num_kv_heads: int = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_key_value_groups: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.kv_repeat: int = self.num_key_value_groups

        self.q_head_size: int = self.num_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        # Adaptive k_rank
        comp_config = getattr(cfg, "compression_config", None)
        if k_rank is None or k_rank == 64:
            effective_rank = getattr(
                comp_config,
                "k_rank",
                getattr(
                    comp_config,
                    "soa_rank",
                    rank if rank is not None else max(128, self.hidden_size // 4),
                ),
            )
        else:
            effective_rank = min(k_rank, self.hidden_size)
        self.k_rank: int = effective_rank
        self.rank = effective_rank
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha

        # Decoder / Causal flag
        _mt = str(getattr(cfg, "model_type", "") or "").lower() if cfg else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_dec = (
            bool(_cfg_get(cfg, "is_decoder", False))
            or bool(_cfg_get(cfg, "is_causal", False))
            or any(_mt.startswith(p) for p in _decoder_prefixes)
            or (src_attn is not None and hasattr(src_attn, "c_attn"))
            or (src_attn is not None and getattr(src_attn, "is_causal", False))
            or (src_attn is not None and getattr(src_attn, "is_decoder", False))
            or bool(kwargs.get("is_causal", False))
            or bool(kwargs.get("is_decoder", False))
        )
        self.is_decoder: bool = is_dec
        self.is_causal: bool = is_dec
        self.is_nested_encoder: bool = False

        # 2. Factorized SVD linear projections for Q, K, V
        self.q_down = nn.Linear(self.hidden_size, self.k_rank, bias=False)
        self.q_up = FactorizedLoRAUp(
            self.k_rank,
            self.q_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        self.k_down = nn.Linear(self.hidden_size, self.k_rank, bias=False)
        self.k_up = FactorizedLoRAUp(
            self.k_rank,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        self.v_down = nn.Linear(self.hidden_size, self.k_rank, bias=False)
        self.v_up = FactorizedLoRAUp(
            self.k_rank,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
        )

        # Output projection with trainable LoRA adapter
        self.out_proj = LoRALinear(
            self.q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = (
            getattr(cfg, "dropout", getattr(cfg, "attention_probs_dropout_prob", dropout))
            if cfg is not None
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and src_attn is not None:
            target_dtype = next(src_attn.parameters(), torch.tensor([])).dtype
            if not isinstance(target_dtype, torch.dtype):
                target_dtype = torch.float32

        # 3. Pretrained weight transfer
        if src_attn is not None:
            self.copy_weights(src_attn)

        if target_dtype is not None:
            self.to(target_dtype)

    @property
    def q_proj(self) -> FactorizedLinear:
        return FactorizedLinear(self.q_down, self.q_up)

    @property
    def k_proj(self) -> FactorizedLinear:
        return FactorizedLinear(self.k_down, self.k_up)

    @property
    def v_proj(self) -> FactorizedLinear:
        return FactorizedLinear(self.v_down, self.v_up)

    @property
    def query(self) -> FactorizedLinear:
        return self.q_proj

    @property
    def key(self) -> FactorizedLinear:
        return self.k_proj

    @property
    def value(self) -> FactorizedLinear:
        return self.v_proj

    def copy_weights(self, old_attn: nn.Module) -> None:
        """Robust weight extraction and SVD factorization across all architectures."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Check nested encoder container
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_out_lin = (
                hasattr(old_attn, "out_lin")
                or hasattr(old_attn, "c_proj")
                or hasattr(old_attn, "o_proj")
                or hasattr(old_attn, "dense")
            )
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query") and not has_out_lin
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                self.is_decoder = True
                self.is_causal = True

                c_attn_mod = src.c_attn
                base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
                c_attn_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)).detach()
                c_attn_b = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
                c_attn_b = c_attn_b.detach() if c_attn_b is not None else None

                c_cls = getattr(base_c, "__class__", type(base_c)).__name__
                if (
                    ("Conv1D" in c_cls)
                    or hasattr(base_c, "nf")
                    or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                    or hasattr(c_attn_mod, "nf")
                    or (c_attn_w.dim() == 2 and c_attn_w.shape[0] < c_attn_w.shape[1])
                ):
                    c_attn_w = c_attn_w.t().contiguous()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size

                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

                c_proj_mod = src.c_proj
                base_p = getattr(c_proj_mod, "base_layer", c_proj_mod)
                out_w = getattr(base_p, "weight", getattr(c_proj_mod, "weight", None)).detach()
                b_o = getattr(base_p, "bias", getattr(c_proj_mod, "bias", None))
                b_o = b_o.detach() if b_o is not None else None

                p_cls = getattr(base_p, "__class__", type(base_p)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(base_p, "nf")
                    or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                    or hasattr(c_proj_mod, "nf")
                    or (out_w.dim() == 2 and out_w.shape[0] < out_w.shape[1])
                ):
                    out_w = out_w.t().contiguous()
                w_o = out_w

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None

                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None

                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # Pattern 5: UniversalAdapter fallback
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    w_q, b_q = adapter.get_weight_bias(src, "q")
                    w_k, b_k = adapter.get_weight_bias(src, "k")
                    w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception:
                    pass

            # Strict validation
            missing = []
            if w_q is None:
                missing.append("Q projection")
            if w_k is None:
                missing.append("K projection")
            if w_v is None:
                missing.append("V projection")
            if w_o is None and not self.is_nested_encoder:
                missing.append("Out projection")

            if missing:
                raise RuntimeError(
                    f"SOAAttention: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
                )

            # SVD Factorization helper
            def _factorize_svd(
                W: torch.Tensor,
                down: nn.Linear,
                up: nn.Linear,
                bias: torch.Tensor | None,
            ) -> None:
                W_f = W.float()
                U, S, Vt = torch.linalg.svd(W_f, full_matrices=False)
                S = S.clamp(min=1e-6)
                r = min(self.k_rank, U.size(1), Vt.size(0))

                down.weight.zero_()
                up.weight.zero_()
                sqrt_s = torch.sqrt(S[:r])
                down.weight[:r, :].copy_((sqrt_s.unsqueeze(1) * Vt[:r, :]).to(down.weight.dtype))
                up.weight[:, :r].copy_((U[:, :r] * sqrt_s).to(up.weight.dtype))

                if bias is not None and up.bias is not None:
                    up.bias.copy_(bias.to(up.bias.dtype))
                elif up.bias is not None:
                    up.bias.zero_()

                # Frobenius check
                rec = up.weight @ down.weight
                f_norm = float(torch.linalg.norm(rec))
                if f_norm < 1e-4:
                    logger.warning(
                        f"SOAAttention: SVD reconstructed weight norm is near zero: {f_norm}"
                    )

            _factorize_svd(w_q, self.q_down, self.q_up, b_q)
            _factorize_svd(w_k, self.k_down, self.k_up, b_k)
            _factorize_svd(w_v, self.v_down, self.v_up, b_v)

            if w_o is not None and not self.is_nested_encoder:
                self.out_proj.weight.copy_(w_o.to(self.out_proj.weight.dtype))
                if b_o is not None and self.out_proj.bias is not None:
                    self.out_proj.bias.copy_(b_o.to(self.out_proj.bias.dtype))
                f_norm_o = float(torch.linalg.norm(self.out_proj.weight))
                if f_norm_o < 1e-4:
                    logger.warning(f"SOAAttention: Out projection norm is near zero: {f_norm_o}")
            elif self.is_nested_encoder:
                if hasattr(self.out_proj, "weight") and self.out_proj.weight.ndim == 2:
                    nn.init.eye_(self.out_proj.weight)
                if hasattr(self.out_proj, "bias") and self.out_proj.bias is not None:
                    nn.init.zeros_(self.out_proj.bias)

            # Freeze base weights
            for mod in [self.q_down, self.k_down, self.v_down]:
                mod.weight.requires_grad = False
                if getattr(mod, "bias", None) is not None:
                    mod.bias.requires_grad = False
            for mod in [self.q_up, self.k_up, self.v_up]:
                if hasattr(mod, "freeze_base_weights"):
                    mod.freeze_base_weights()
                else:
                    mod.weight.requires_grad = False
                    if getattr(mod, "bias", None) is not None:
                        mod.bias.requires_grad = False
                    if getattr(mod, "lora_A", None) is not None:
                        mod.lora_A.requires_grad = True
                    if getattr(mod, "lora_B", None) is not None:
                        mod.lora_B.requires_grad = True
            if hasattr(self.out_proj, "freeze_base_weights"):
                self.out_proj.freeze_base_weights()
            else:
                self.out_proj.weight.requires_grad = False
                if getattr(self.out_proj, "bias", None) is not None:
                    self.out_proj.bias.requires_grad = False
                if getattr(self.out_proj, "lora_A", None) is not None:
                    self.out_proj.lora_A.requires_grad = True
                if getattr(self.out_proj, "lora_B", None) is not None:
                    self.out_proj.lora_B.requires_grad = True

    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> "SOAAttention":
        return cls(model_config=model_config, pretrained_attn=old_attn, **kwargs)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        B, N, _ = hidden_states.shape

        # Factorized projections
        z_q = self.q_down(hidden_states)
        z_k = self.k_down(hidden_states)
        z_v = self.v_down(hidden_states)

        q = self.q_up(z_q)
        k = self.k_up(z_k)
        v = self.v_up(z_v)

        H = self.num_heads
        H_kv = self.num_kv_heads
        D = self.head_dim

        q = q.view(B, N, H, D).transpose(1, 2).contiguous()
        k = k.view(B, N, H_kv, D).transpose(1, 2).contiguous()
        v = v.view(B, N, H_kv, D).transpose(1, 2).contiguous()

        # Partial Rotary Position Embeddings (RoPE)
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
            cos, sin = pos_emb
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # Grouped-Query Attention expansion
        if H_kv != H:
            k = repeat_kv(k, self.num_key_value_groups)
            v = repeat_kv(v, self.num_key_value_groups)

        scale = 1.0 / math.sqrt(D)
        # Memory-efficient: use SDPA to avoid OOM on large seqs

        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

            _sdpa_out = torch.nn.functional.scaled_dot_product_attention(
                _qs,
                _ks,
                _vs,
                is_causal=getattr(self, "is_decoder", False) and N > 1,
                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, self.all_head_size)
            )

            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None

        except Exception:
            pass

        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        # Causal mask for decoders
        is_causal = (
            self.is_decoder or getattr(self, "is_causal", False) or kwargs.get("is_causal", False)
        ) and (encoder_hidden_states is None)
        if is_causal and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full((N, N), float("-inf"), device=scores.device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        # Standard Attention Mask
        mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask")
        if mask is not None:
            if mask.dim() == 4 and mask.dtype == torch.bool:
                scores = scores.masked_fill(~mask, -10000.0)
            else:
                scores = apply_attention_mask(scores, mask)

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights.to(v.dtype))

        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)
            attn_weights = attn_weights * head_mask

        context = torch.matmul(attn_weights, v)
        context = context.transpose(1, 2).contiguous().view(B, N, H * D)

        if self.is_nested_encoder:
            return (context, attn_weights if output_attentions else None)

        output = self.out_proj(context)
        return (output, attn_weights if output_attentions else None)

Shared Orthogonal Attention (SOA) with Per-Projection Latent Bottlenecks.

Factorizes dense Q, K, and V linear projections into low-rank bottleneck pairs initialized via truncated SVD from pretrained transformer weights.

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_gqa : bool

The type of the None singleton.

Static methods

def from_distilbert_pretrained(old_attn: torch.nn.modules.module.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> SOAAttention

Instance variables

prop k_proj : FactorizedLinear

Expand source code

@property
def k_proj(self) -> FactorizedLinear:
    return FactorizedLinear(self.k_down, self.k_up)

prop key : FactorizedLinear

Expand source code

@property
def key(self) -> FactorizedLinear:
    return self.k_proj

prop q_proj : FactorizedLinear

Expand source code

@property
def q_proj(self) -> FactorizedLinear:
    return FactorizedLinear(self.q_down, self.q_up)

prop query : FactorizedLinear

Expand source code

@property
def query(self) -> FactorizedLinear:
    return self.q_proj

prop v_proj : FactorizedLinear

Expand source code

@property
def v_proj(self) -> FactorizedLinear:
    return FactorizedLinear(self.v_down, self.v_up)

prop value : FactorizedLinear

Expand source code

@property
def value(self) -> FactorizedLinear:
    return self.v_proj

Methods

def copy_weights(self, old_attn: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    """Robust weight extraction and SVD factorization across all architectures."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Check nested encoder container
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        has_out_lin = (
            hasattr(old_attn, "out_lin")
            or hasattr(old_attn, "c_proj")
            or hasattr(old_attn, "o_proj")
            or hasattr(old_attn, "dense")
        )
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query") and not has_out_lin
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight
                b_o = old_attn.output.dense.bias
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            self.is_decoder = True
            self.is_causal = True

            c_attn_mod = src.c_attn
            base_c = getattr(c_attn_mod, "base_layer", c_attn_mod)
            c_attn_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)).detach()
            c_attn_b = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None))
            c_attn_b = c_attn_b.detach() if c_attn_b is not None else None

            c_cls = getattr(base_c, "__class__", type(base_c)).__name__
            if (
                ("Conv1D" in c_cls)
                or hasattr(base_c, "nf")
                or ("Conv1D" in getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__)
                or hasattr(c_attn_mod, "nf")
                or (c_attn_w.dim() == 2 and c_attn_w.shape[0] < c_attn_w.shape[1])
            ):
                c_attn_w = c_attn_w.t().contiguous()

            q_dim = self.q_head_size
            k_dim = self.kv_head_size
            v_dim = self.kv_head_size

            w_q = c_attn_w[:q_dim, :]
            w_k = c_attn_w[q_dim : q_dim + k_dim, :]
            w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

            if c_attn_b is not None:
                b_q = c_attn_b[:q_dim]
                b_k = c_attn_b[q_dim : q_dim + k_dim]
                b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

            c_proj_mod = src.c_proj
            base_p = getattr(c_proj_mod, "base_layer", c_proj_mod)
            out_w = getattr(base_p, "weight", getattr(c_proj_mod, "weight", None)).detach()
            b_o = getattr(base_p, "bias", getattr(c_proj_mod, "bias", None))
            b_o = b_o.detach() if b_o is not None else None

            p_cls = getattr(base_p, "__class__", type(base_p)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(base_p, "nf")
                or ("Conv1D" in getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__)
                or hasattr(c_proj_mod, "nf")
                or (out_w.dim() == 2 and out_w.shape[0] < out_w.shape[1])
            ):
                out_w = out_w.t().contiguous()
            w_o = out_w

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None

            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None

            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # Pattern 5: UniversalAdapter fallback
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                w_q, b_q = adapter.get_weight_bias(src, "q")
                w_k, b_k = adapter.get_weight_bias(src, "k")
                w_v, b_v = adapter.get_weight_bias(src, "v")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(src, "o")
            except Exception:
                pass

        # Strict validation
        missing = []
        if w_q is None:
            missing.append("Q projection")
        if w_k is None:
            missing.append("K projection")
        if w_v is None:
            missing.append("V projection")
        if w_o is None and not self.is_nested_encoder:
            missing.append("Out projection")

        if missing:
            raise RuntimeError(
                f"SOAAttention: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
            )

        # SVD Factorization helper
        def _factorize_svd(
            W: torch.Tensor,
            down: nn.Linear,
            up: nn.Linear,
            bias: torch.Tensor | None,
        ) -> None:
            W_f = W.float()
            U, S, Vt = torch.linalg.svd(W_f, full_matrices=False)
            S = S.clamp(min=1e-6)
            r = min(self.k_rank, U.size(1), Vt.size(0))

            down.weight.zero_()
            up.weight.zero_()
            sqrt_s = torch.sqrt(S[:r])
            down.weight[:r, :].copy_((sqrt_s.unsqueeze(1) * Vt[:r, :]).to(down.weight.dtype))
            up.weight[:, :r].copy_((U[:, :r] * sqrt_s).to(up.weight.dtype))

            if bias is not None and up.bias is not None:
                up.bias.copy_(bias.to(up.bias.dtype))
            elif up.bias is not None:
                up.bias.zero_()

            # Frobenius check
            rec = up.weight @ down.weight
            f_norm = float(torch.linalg.norm(rec))
            if f_norm < 1e-4:
                logger.warning(
                    f"SOAAttention: SVD reconstructed weight norm is near zero: {f_norm}"
                )

        _factorize_svd(w_q, self.q_down, self.q_up, b_q)
        _factorize_svd(w_k, self.k_down, self.k_up, b_k)
        _factorize_svd(w_v, self.v_down, self.v_up, b_v)

        if w_o is not None and not self.is_nested_encoder:
            self.out_proj.weight.copy_(w_o.to(self.out_proj.weight.dtype))
            if b_o is not None and self.out_proj.bias is not None:
                self.out_proj.bias.copy_(b_o.to(self.out_proj.bias.dtype))
            f_norm_o = float(torch.linalg.norm(self.out_proj.weight))
            if f_norm_o < 1e-4:
                logger.warning(f"SOAAttention: Out projection norm is near zero: {f_norm_o}")
        elif self.is_nested_encoder:
            if hasattr(self.out_proj, "weight") and self.out_proj.weight.ndim == 2:
                nn.init.eye_(self.out_proj.weight)
            if hasattr(self.out_proj, "bias") and self.out_proj.bias is not None:
                nn.init.zeros_(self.out_proj.bias)

        # Freeze base weights
        for mod in [self.q_down, self.k_down, self.v_down]:
            mod.weight.requires_grad = False
            if getattr(mod, "bias", None) is not None:
                mod.bias.requires_grad = False
        for mod in [self.q_up, self.k_up, self.v_up]:
            if hasattr(mod, "freeze_base_weights"):
                mod.freeze_base_weights()
            else:
                mod.weight.requires_grad = False
                if getattr(mod, "bias", None) is not None:
                    mod.bias.requires_grad = False
                if getattr(mod, "lora_A", None) is not None:
                    mod.lora_A.requires_grad = True
                if getattr(mod, "lora_B", None) is not None:
                    mod.lora_B.requires_grad = True
        if hasattr(self.out_proj, "freeze_base_weights"):
            self.out_proj.freeze_base_weights()
        else:
            self.out_proj.weight.requires_grad = False
            if getattr(self.out_proj, "bias", None) is not None:
                self.out_proj.bias.requires_grad = False
            if getattr(self.out_proj, "lora_A", None) is not None:
                self.out_proj.lora_A.requires_grad = True
            if getattr(self.out_proj, "lora_B", None) is not None:
                self.out_proj.lora_B.requires_grad = True

Robust weight extraction and SVD factorization across all architectures.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class SVDCompressedAttention (config: Any | None = None,
old_attn: Any | None = None,
dropout: float = 0.1,
layer_idx: int = 0,
ranks: dict[str, int] | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
add_sparse: bool = False,
sparsity_level: float = 0.7,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("svd_compressed", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_compressed_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_compressed_module", supports_causal=True, supports_gqa=True)
class SVDCompressedAttention(BaseAttention):
    r"""Multi-head attention with truncated SVD projections, trainable LoRA adapters,
    GQA support, RoPE support, and optional sparse residuals.
    """

    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        old_attn: Any | None = None,
        dropout: float = 0.1,
        layer_idx: int = 0,
        ranks: dict[str, int] | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        add_sparse: bool = False,
        sparsity_level: float = 0.7,
        **kwargs: Any,
    ) -> None:
        model_config = config or kwargs.get("model_config")
        super().__init__(model_config)
        # 1. Dimensions & Architecture
        self.hidden_size: int = int(
            _cfg_get(model_config, "hidden_size", kwargs.get("hidden_size", 768))
        )
        self.num_heads: int = int(
            _cfg_get(
                model_config,
                "num_attention_heads",
                _cfg_get(model_config, "num_heads", kwargs.get("num_attention_heads", 12)),
            )
        )
        self.head_dim: int = int(
            _cfg_get(
                model_config,
                "head_dim",
                kwargs.get(
                    "head_dim",
                    self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
                ),
            )
        )

        cfg_kv = _cfg_get(
            model_config,
            "num_key_value_heads",
            _cfg_get(model_config, "num_kv_heads", None),
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = max(1, self.num_heads // self.num_kv_heads)
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.q_head_size: int = self.num_heads * self.head_dim
        self.all_head_size: int = self.q_head_size

        # 2. Causal decoder flag detection
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(_cfg_get(model_config, "model_type", "") or "").lower()
            self.is_decoder = (
                bool(_cfg_get(model_config, "is_decoder", False))
                or any(
                    _mt.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "falcon",
                        "bloom",
                        "opt",
                    )
                )
                or (old_attn is not None and hasattr(old_attn, "c_attn"))
            )

        self.is_nested_encoder: bool = False
        default_r = int(_cfg_get(model_config, "svd_rank", _cfg_get(model_config, "rank", 64)))
        default_r = max(1, min(default_r, self.hidden_size, self.all_head_size))
        self.ranks: dict[str, int] = {}
        for k in ("q", "k", "v", "o"):
            val = ranks.get(k, default_r) if ranks else default_r
            self.ranks[k] = max(1, min(int(val), self.hidden_size, self.all_head_size))

        self.add_sparse: bool = add_sparse
        self.sparsity_level: float = sparsity_level
        self.lora_rank = max(
            1,
            min(
                int(lora_rank if lora_rank is not None else 16),
                self.hidden_size,
                self.all_head_size,
            ),
        )
        self.lora_alpha = lora_alpha
        self.scaling = lora_alpha / max(1, self.lora_rank)

        # 3. Factorized projections
        self.q_down = nn.Linear(self.hidden_size, self.ranks["q"], bias=False)
        self.q_up = nn.Linear(self.ranks["q"], self.q_head_size, bias=True)
        self.k_down = nn.Linear(self.hidden_size, self.ranks["k"], bias=False)
        self.k_up = nn.Linear(self.ranks["k"], self.kv_head_size, bias=True)
        self.v_down = nn.Linear(self.hidden_size, self.ranks["v"], bias=False)
        self.v_up = nn.Linear(self.ranks["v"], self.kv_head_size, bias=True)
        self.o_down = nn.Linear(self.q_head_size, self.ranks["o"], bias=False)
        self.o_up = nn.Linear(self.ranks["o"], self.hidden_size, bias=True)

        for mod in [
            self.q_down,
            self.q_up,
            self.k_down,
            self.k_up,
            self.v_down,
            self.v_up,
            self.o_down,
            self.o_up,
        ]:
            if isinstance(mod, nn.Linear):
                nn.init.orthogonal_(mod.weight)
                if mod.bias is not None:
                    nn.init.zeros_(mod.bias)
                mod.weight.requires_grad = False
                if mod.bias is not None:
                    mod.bias.requires_grad = False

        # 4. Trainable LoRA parameters
        if self.lora_rank > 0:
            self.q_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.q_lora_B = nn.Parameter(torch.empty(self.q_head_size, self.lora_rank))
            self.k_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.k_lora_B = nn.Parameter(torch.empty(self.kv_head_size, self.lora_rank))
            self.v_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.v_lora_B = nn.Parameter(torch.empty(self.kv_head_size, self.lora_rank))
            self.o_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.q_head_size))
            self.o_lora_B = nn.Parameter(torch.empty(self.hidden_size, self.lora_rank))

            for p_a in (self.q_lora_A, self.k_lora_A, self.v_lora_A, self.o_lora_A):
                nn.init.kaiming_uniform_(p_a, a=math.sqrt(5))
            for p_b in (self.q_lora_B, self.k_lora_B, self.v_lora_B, self.o_lora_B):
                nn.init.zeros_(p_b)
        else:
            self.register_parameter("q_lora_A", None)
            self.register_parameter("q_lora_B", None)
            self.register_parameter("k_lora_A", None)
            self.register_parameter("k_lora_B", None)
            self.register_parameter("v_lora_A", None)
            self.register_parameter("v_lora_B", None)
            self.register_parameter("o_lora_A", None)
            self.register_parameter("o_lora_B", None)

        if add_sparse:
            self.q_sparse_weight = nn.Parameter(torch.zeros(self.q_head_size, self.hidden_size))
            self.k_sparse_weight = nn.Parameter(torch.zeros(self.kv_head_size, self.hidden_size))
            self.v_sparse_weight = nn.Parameter(torch.zeros(self.kv_head_size, self.hidden_size))
            self.o_sparse_weight = nn.Parameter(torch.zeros(self.hidden_size, self.q_head_size))
        else:
            self.register_parameter("q_sparse_weight", None)
            self.register_parameter("k_sparse_weight", None)
            self.register_parameter("v_sparse_weight", None)
            self.register_parameter("o_sparse_weight", None)

        dropout_p = _cfg_get(
            model_config,
            "dropout",
            _cfg_get(
                model_config,
                "hidden_dropout_prob",
                _cfg_get(model_config, "attention_probs_dropout_prob", dropout),
            ),
        )
        self.dropout = nn.Dropout(float(dropout_p) if dropout_p is not None else 0.1)

        src_attn = old_attn or kwargs.get("pretrained_attn")
        src_p = next(src_attn.parameters(), None) if src_attn is not None else None
        target_dtype = kwargs.get("dtype", src_p.dtype if src_p is not None else torch.float32)
        target_device = kwargs.get("device", src_p.device if src_p is not None else None)

        if src_attn is not None:
            self.copy_weights(src_attn)

        if target_device is not None and target_dtype is not None:
            self.to(device=target_device, dtype=target_dtype)
        elif target_device is not None:
            self.to(device=target_device)
        elif target_dtype is not None:
            self.to(target_dtype)

    def _set_factorised_weights(
        self,
        W_tensor: torch.Tensor,
        bias_tensor: torch.Tensor | None,
        down: nn.Linear,
        up: nn.Linear,
        rank: int,
    ) -> None:
        with torch.no_grad():
            W = W_tensor.float()
            d_out, d_in = W.shape
            r = min(rank, d_out, d_in, down.out_features, up.in_features)
            r = max(r, 1)

            U, s, VT = torch.linalg.svd(W, full_matrices=False)
            sqrt_s = torch.sqrt(torch.clamp(s[:r], min=1e-12))
            down_weight = (sqrt_s.unsqueeze(1) * VT[:r, :]).to(
                device=down.weight.device, dtype=down.weight.dtype
            )
            up_weight = (U[:, :r] * sqrt_s.unsqueeze(0)).to(
                device=up.weight.device, dtype=up.weight.dtype
            )

            down.weight.data.zero_()
            up.weight.data.zero_()
            down.weight.data[:r, :].copy_(down_weight)
            up.weight.data[:, :r].copy_(up_weight)
            if bias_tensor is not None and up.bias is not None:
                up.bias.data.copy_(bias_tensor.to(device=up.bias.device, dtype=up.bias.dtype))

    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Detect nested encoder container or inner self-attention
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_out_proj = any(
                hasattr(old_attn, attr)
                for attr in (
                    "dense",
                    "out_proj",
                    "c_proj",
                    "out_lin",
                    "o_proj",
                    "wo",
                    "o",
                    "proj",
                )
            )
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query") and not has_out_proj
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o, b_o = _unpack_weight_bias(old_attn.output.dense)
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_w, c_b = _unpack_weight_bias(src.c_attn)
                if c_w is not None:
                    if c_w.shape[0] == self.hidden_size and c_w.shape[1] != self.hidden_size:
                        c_w = c_w.t().contiguous()

                    q_dim = self.q_head_size
                    k_dim = self.kv_head_size
                    v_dim = self.kv_head_size

                    w_q = c_w[:q_dim, :]
                    w_k = c_w[q_dim : q_dim + k_dim, :]
                    w_v = c_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                    if c_b is not None:
                        b_q = c_b[:q_dim]
                        b_k = c_b[q_dim : q_dim + k_dim]
                        b_v = c_b[q_dim + k_dim : q_dim + k_dim + v_dim]

                out_w, out_b = _unpack_weight_bias(src.c_proj)
                if out_w is not None:
                    if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                        out_w = out_w.t().contiguous()
                    w_o = out_w
                    b_o = out_b

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q, b_q = _unpack_weight_bias(src.query)
                w_k, b_k = _unpack_weight_bias(src.key)
                w_v, b_v = _unpack_weight_bias(src.value)

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q, b_q = _unpack_weight_bias(src.q_proj)
                w_k, b_k = _unpack_weight_bias(src.k_proj)
                w_v, b_v = _unpack_weight_bias(src.v_proj)
                if hasattr(src, "o_proj"):
                    w_o, b_o = _unpack_weight_bias(src.o_proj)

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q, b_q = _unpack_weight_bias(src.q_lin)
                w_k, b_k = _unpack_weight_bias(src.k_lin)
                w_v, b_v = _unpack_weight_bias(src.v_lin)
                if hasattr(src, "out_lin"):
                    w_o, b_o = _unpack_weight_bias(src.out_lin)

            # Fallback: UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter(old_attn)
                    extracted = adapter.extract_qkv_out()
                    if extracted:
                        w_q = extracted.get("q_weight")
                        b_q = extracted.get("q_bias")
                        w_k = extracted.get("k_weight")
                        b_k = extracted.get("k_bias")
                        w_v = extracted.get("v_weight")
                        b_v = extracted.get("v_bias")
                        if w_o is None:
                            w_o = extracted.get("out_weight")
                            b_o = extracted.get("out_bias")
                except Exception as e:
                    log.debug("UniversalAdapter fallback failed: %s", e)

            # Verification of essential projections
            missing = []
            if w_q is None:
                missing.append("query/q_proj")
            if w_k is None:
                missing.append("key/k_proj")
            if w_v is None:
                missing.append("value/v_proj")
            if missing:
                log.warning("Failed to extract required projections: %s", missing)
                return

            # Frobenius norm checks
            for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    log.warning(f"Copied weight {name} has near-zero Frobenius norm (< 1e-4)")

            # Orient weights if needed
            def _orient(w: torch.Tensor, out_dim: int, in_dim: int) -> torch.Tensor:
                if w.shape == (out_dim, in_dim):
                    return w
                if w.shape == (in_dim, out_dim):
                    return w.t().contiguous()
                return w

            w_q = _orient(w_q, self.q_head_size, self.hidden_size)
            # Factorize and set Q
            self._set_factorised_weights(w_q, b_q, self.q_down, self.q_up, self.ranks["q"])

            # Factorize and set K & V with GQA pooling if needed
            if w_k.shape[0] == self.kv_head_size:
                w_k = _orient(w_k, self.kv_head_size, self.hidden_size)
                self._set_factorised_weights(w_k, b_k, self.k_down, self.k_up, self.ranks["k"])
            else:
                rep = max(1, w_k.shape[0] // self.kv_head_size)
                w_k_pooled = (
                    w_k.view(self.num_kv_heads, rep, self.head_dim, -1)
                    .mean(dim=1)
                    .view(self.kv_head_size, -1)
                )
                b_k_pooled = (
                    b_k.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                    if b_k is not None
                    else None
                )
                self._set_factorised_weights(
                    w_k_pooled, b_k_pooled, self.k_down, self.k_up, self.ranks["k"]
                )

            if w_v.shape[0] == self.kv_head_size:
                w_v = _orient(w_v, self.kv_head_size, self.hidden_size)
                self._set_factorised_weights(w_v, b_v, self.v_down, self.v_up, self.ranks["v"])
            else:
                rep = max(1, w_v.shape[0] // self.kv_head_size)
                w_v_pooled = (
                    w_v.view(self.num_kv_heads, rep, self.head_dim, -1)
                    .mean(dim=1)
                    .view(self.kv_head_size, -1)
                )
                b_v_pooled = (
                    b_v.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                    if b_v is not None
                    else None
                )
                self._set_factorised_weights(
                    w_v_pooled, b_v_pooled, self.v_down, self.v_up, self.ranks["v"]
                )

            # Set Out projection if not nested encoder
            if w_o is not None and not self.is_nested_encoder:
                if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                    log.warning("Copied out_proj weight has near-zero Frobenius norm (< 1e-4)")
                w_o = _orient(w_o, self.hidden_size, self.q_head_size)
                self._set_factorised_weights(w_o, b_o, self.o_down, self.o_up, self.ranks["o"])

            # Initialize sparse residuals if enabled
            if self.add_sparse:
                low_rank_q = self.q_up.weight @ self.q_down.weight
                res_q = w_q - low_rank_q
                self.q_sparse_weight.data.copy_(res_q)

                low_rank_k = self.k_up.weight @ self.k_down.weight
                res_k = (w_k if w_k.shape[0] == self.kv_head_size else w_k_pooled) - low_rank_k
                self.k_sparse_weight.data.copy_(res_k)

                low_rank_v = self.v_up.weight @ self.v_down.weight
                res_v = (w_v if w_v.shape[0] == self.kv_head_size else w_v_pooled) - low_rank_v
                self.v_sparse_weight.data.copy_(res_v)

                if w_o is not None and not self.is_nested_encoder:
                    low_rank_o = self.o_up.weight @ self.o_down.weight
                    res_o = w_o - low_rank_o
                    self.o_sparse_weight.data.copy_(res_o)

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        # Defensive device alignment
        if self.q_down.weight.device != hidden_states.device:
            self.to(device=hidden_states.device)

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        kv_len = kv_input.shape[1]

        def proj_q(x: torch.Tensor) -> torch.Tensor:
            if x.device != self.q_down.weight.device:
                x = x.to(self.q_down.weight.device)
            out = self.q_up(self.q_down(x))
            if self.lora_rank > 0 and self.q_lora_A is not None and self.q_lora_B is not None:
                out = out + F.linear(F.linear(x, self.q_lora_A), self.q_lora_B) * self.scaling
            if self.add_sparse and self.q_sparse_weight is not None:
                out = out + F.linear(x, self.q_sparse_weight)
            return out

        def proj_k(x: torch.Tensor) -> torch.Tensor:
            if x.device != self.k_down.weight.device:
                x = x.to(self.k_down.weight.device)
            out = self.k_up(self.k_down(x))
            if self.lora_rank > 0 and self.k_lora_A is not None and self.k_lora_B is not None:
                out = out + F.linear(F.linear(x, self.k_lora_A), self.k_lora_B) * self.scaling
            if self.add_sparse and self.k_sparse_weight is not None:
                out = out + F.linear(x, self.k_sparse_weight)
            return out

        def proj_v(x: torch.Tensor) -> torch.Tensor:
            if x.device != self.v_down.weight.device:
                x = x.to(self.v_down.weight.device)
            out = self.v_up(self.v_down(x))
            if self.lora_rank > 0 and self.v_lora_A is not None and self.v_lora_B is not None:
                out = out + F.linear(F.linear(x, self.v_lora_A), self.v_lora_B) * self.scaling
            if self.add_sparse and self.v_sparse_weight is not None:
                out = out + F.linear(x, self.v_sparse_weight)
            return out

        def proj_o(x: torch.Tensor) -> torch.Tensor:
            if x.device != self.o_down.weight.device:
                x = x.to(self.o_down.weight.device)
            out = self.o_up(self.o_down(x))
            if self.lora_rank > 0 and self.o_lora_A is not None and self.o_lora_B is not None:
                out = out + F.linear(F.linear(x, self.o_lora_A), self.o_lora_B) * self.scaling
            if self.add_sparse and self.o_sparse_weight is not None:
                out = out + F.linear(x, self.o_sparse_weight)
            return out

        # 1. Projections
        Q = (
            proj_q(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            proj_k(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V = (
            proj_v(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. RoPE support
        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
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # 3. GQA expansion
        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        # 4. Attention scores in FP32
        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(Q.float(), K.float().transpose(-1, -2).contiguous()) * scale

        # 5. Masking
        if self.is_decoder and not is_cross and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full(
                        (N, kv_len),
                        float("-inf"),
                        device=scores.device,
                        dtype=torch.float32,
                    ),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            scores = apply_attention_mask(scores, eff_mask)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        attn_weights = self.dropout(attn_weights)

        # 6. Safe head_mask application
        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn_weights = attn_weights * head_mask

        # 7. Context computation & Head Merge
        context = (
            torch.matmul(attn_weights, V)
            .transpose(1, 2)
            .contiguous()
            .view(B, N, self.all_head_size)
        )
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # 8. Output projection (bypassed if nested encoder)
        if self.is_nested_encoder:
            output = context
        else:
            output = proj_o(context)

        if output.shape[-1] != hidden_states.shape[-1]:
            if output.shape[-1] < hidden_states.shape[-1]:
                output = F.pad(output, (0, hidden_states.shape[-1] - output.shape[-1]))
            else:
                output = output[..., : hidden_states.shape[-1]]

        return output, (attn_weights if output_attentions else None)

Multi-head attention with truncated SVD projections, trainable LoRA adapters, GQA support, RoPE support, and optional sparse residuals.

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_gqa : bool

The type of the None singleton.

Methods

def copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Detect nested encoder container or inner self-attention
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        has_out_proj = any(
            hasattr(old_attn, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query") and not has_out_proj
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o, b_o = _unpack_weight_bias(old_attn.output.dense)
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_w, c_b = _unpack_weight_bias(src.c_attn)
            if c_w is not None:
                if c_w.shape[0] == self.hidden_size and c_w.shape[1] != self.hidden_size:
                    c_w = c_w.t().contiguous()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size

                w_q = c_w[:q_dim, :]
                w_k = c_w[q_dim : q_dim + k_dim, :]
                w_v = c_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                if c_b is not None:
                    b_q = c_b[:q_dim]
                    b_k = c_b[q_dim : q_dim + k_dim]
                    b_v = c_b[q_dim + k_dim : q_dim + k_dim + v_dim]

            out_w, out_b = _unpack_weight_bias(src.c_proj)
            if out_w is not None:
                if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = out_b

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q, b_q = _unpack_weight_bias(src.query)
            w_k, b_k = _unpack_weight_bias(src.key)
            w_v, b_v = _unpack_weight_bias(src.value)

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q, b_q = _unpack_weight_bias(src.q_proj)
            w_k, b_k = _unpack_weight_bias(src.k_proj)
            w_v, b_v = _unpack_weight_bias(src.v_proj)
            if hasattr(src, "o_proj"):
                w_o, b_o = _unpack_weight_bias(src.o_proj)

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q, b_q = _unpack_weight_bias(src.q_lin)
            w_k, b_k = _unpack_weight_bias(src.k_lin)
            w_v, b_v = _unpack_weight_bias(src.v_lin)
            if hasattr(src, "out_lin"):
                w_o, b_o = _unpack_weight_bias(src.out_lin)

        # Fallback: UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter(old_attn)
                extracted = adapter.extract_qkv_out()
                if extracted:
                    w_q = extracted.get("q_weight")
                    b_q = extracted.get("q_bias")
                    w_k = extracted.get("k_weight")
                    b_k = extracted.get("k_bias")
                    w_v = extracted.get("v_weight")
                    b_v = extracted.get("v_bias")
                    if w_o is None:
                        w_o = extracted.get("out_weight")
                        b_o = extracted.get("out_bias")
            except Exception as e:
                log.debug("UniversalAdapter fallback failed: %s", e)

        # Verification of essential projections
        missing = []
        if w_q is None:
            missing.append("query/q_proj")
        if w_k is None:
            missing.append("key/k_proj")
        if w_v is None:
            missing.append("value/v_proj")
        if missing:
            log.warning("Failed to extract required projections: %s", missing)
            return

        # Frobenius norm checks
        for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                log.warning(f"Copied weight {name} has near-zero Frobenius norm (< 1e-4)")

        # Orient weights if needed
        def _orient(w: torch.Tensor, out_dim: int, in_dim: int) -> torch.Tensor:
            if w.shape == (out_dim, in_dim):
                return w
            if w.shape == (in_dim, out_dim):
                return w.t().contiguous()
            return w

        w_q = _orient(w_q, self.q_head_size, self.hidden_size)
        # Factorize and set Q
        self._set_factorised_weights(w_q, b_q, self.q_down, self.q_up, self.ranks["q"])

        # Factorize and set K & V with GQA pooling if needed
        if w_k.shape[0] == self.kv_head_size:
            w_k = _orient(w_k, self.kv_head_size, self.hidden_size)
            self._set_factorised_weights(w_k, b_k, self.k_down, self.k_up, self.ranks["k"])
        else:
            rep = max(1, w_k.shape[0] // self.kv_head_size)
            w_k_pooled = (
                w_k.view(self.num_kv_heads, rep, self.head_dim, -1)
                .mean(dim=1)
                .view(self.kv_head_size, -1)
            )
            b_k_pooled = (
                b_k.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                if b_k is not None
                else None
            )
            self._set_factorised_weights(
                w_k_pooled, b_k_pooled, self.k_down, self.k_up, self.ranks["k"]
            )

        if w_v.shape[0] == self.kv_head_size:
            w_v = _orient(w_v, self.kv_head_size, self.hidden_size)
            self._set_factorised_weights(w_v, b_v, self.v_down, self.v_up, self.ranks["v"])
        else:
            rep = max(1, w_v.shape[0] // self.kv_head_size)
            w_v_pooled = (
                w_v.view(self.num_kv_heads, rep, self.head_dim, -1)
                .mean(dim=1)
                .view(self.kv_head_size, -1)
            )
            b_v_pooled = (
                b_v.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                if b_v is not None
                else None
            )
            self._set_factorised_weights(
                w_v_pooled, b_v_pooled, self.v_down, self.v_up, self.ranks["v"]
            )

        # Set Out projection if not nested encoder
        if w_o is not None and not self.is_nested_encoder:
            if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                log.warning("Copied out_proj weight has near-zero Frobenius norm (< 1e-4)")
            w_o = _orient(w_o, self.hidden_size, self.q_head_size)
            self._set_factorised_weights(w_o, b_o, self.o_down, self.o_up, self.ranks["o"])

        # Initialize sparse residuals if enabled
        if self.add_sparse:
            low_rank_q = self.q_up.weight @ self.q_down.weight
            res_q = w_q - low_rank_q
            self.q_sparse_weight.data.copy_(res_q)

            low_rank_k = self.k_up.weight @ self.k_down.weight
            res_k = (w_k if w_k.shape[0] == self.kv_head_size else w_k_pooled) - low_rank_k
            self.k_sparse_weight.data.copy_(res_k)

            low_rank_v = self.v_up.weight @ self.v_down.weight
            res_v = (w_v if w_v.shape[0] == self.kv_head_size else w_v_pooled) - low_rank_v
            self.v_sparse_weight.data.copy_(res_v)

            if w_o is not None and not self.is_nested_encoder:
                low_rank_o = self.o_up.weight @ self.o_down.weight
                res_o = w_o - low_rank_o
                self.o_sparse_weight.data.copy_(res_o)

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class SVDLLMV2Attention (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = 192,
energy_threshold: float = 0.999,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
dropout: float = 0.1,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("svd_llm_v2", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_llm_v2_attention")
class SVDLLMV2Attention(BaseAttention):
    r"""SVD-factorized attention layer with explicit U, Sigma, V^H matrix buffers
    and trainable LoRA recovery adapters.
    """

    supports_gqa: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "low_rank"

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = 192,
        energy_threshold: float = 0.999,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        dropout: float = 0.1,
        **kwargs: Any,
    ) -> None:
        comp_config = _cfg_get(config, "compression_config", kwargs.get("comp_config"))
        actual_rank = (
            rank
            if rank is not None
            else _cfg_get(
                comp_config,
                "rank",
                _cfg_get(config, "svd_rank", _cfg_get(config, "rank", 192)),
            )
        )
        super().__init__(config, actual_rank)

        # 1. Dimensions & Architecture
        self.hidden_size: int = kwargs.get(
            "hidden_size",
            _cfg_get(
                config,
                "hidden_size",
                _cfg_get(config, "dim", _cfg_get(config, "d_model", 768)),
            ),
        )
        self.num_heads: int = kwargs.get(
            "num_heads",
            kwargs.get(
                "num_attention_heads",
                _cfg_get(
                    config,
                    "num_attention_heads",
                    _cfg_get(config, "num_heads", _cfg_get(config, "n_heads", 12)),
                ),
            ),
        )

        cfg_head_dim = _cfg_get(config, "head_dim", None)
        adapter_head_dim = (
            getattr(getattr(config, "adapter", None), "head_dim", None) if config else None
        )
        inferred_head_dim = kwargs.get("head_dim", cfg_head_dim or adapter_head_dim)
        if inferred_head_dim is None and old_attn is not None:
            for attr in ("q_proj", "query", "q_lin"):
                if hasattr(old_attn, attr):
                    layer_q = getattr(old_attn, attr)
                    if hasattr(layer_q, "weight"):
                        inferred_head_dim = layer_q.weight.shape[0] // self.num_heads
                        break
        self.head_dim: int = int(
            inferred_head_dim
            if inferred_head_dim is not None
            else (self.hidden_size // self.num_heads if self.num_heads > 0 else 64)
        )

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.q_head_size: int = self.num_heads * self.head_dim
        self.all_head_size: int = self.q_head_size

        self.rank: int = actual_rank
        self.energy_threshold: float = energy_threshold
        self.lora_rank: int = _cfg_get(
            config,
            "lora_rank",
            _cfg_get(comp_config, "lora_rank", lora_rank if lora_rank is not None else 16),
        )
        self.lora_alpha: float = _cfg_get(
            config, "lora_alpha", _cfg_get(comp_config, "lora_alpha", lora_alpha)
        )
        self.scaling: float = self.lora_alpha / max(1, self.lora_rank)

        # 2. Causal decoder flag detection
        _mt = str(_cfg_get(config, "model_type", "") or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "mpt",
        )
        self.is_decoder = (
            bool(_cfg_get(config, "is_decoder", False))
            or bool(_cfg_get(config, "causal", False))
            or any(_mt.startswith(p) for p in _decoder_prefixes)
            or (old_attn is not None and hasattr(old_attn, "c_attn"))
            or bool(kwargs.get("is_causal", False))
        )
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)

        self.is_nested_encoder: bool = False

        # Device & Dtype detection from old_attn
        src_p = next(old_attn.parameters(), None) if old_attn is not None else None
        target_dtype = kwargs.get("dtype", src_p.dtype if src_p is not None else torch.float32)
        target_device = kwargs.get("device", src_p.device if src_p is not None else None)

        # 3. Create parameter buffers for SVD factors
        self.query_u = nn.Parameter(
            torch.empty(self.q_head_size, self.rank, device=target_device, dtype=target_dtype)
        )
        self.query_s = nn.Parameter(
            torch.empty(self.rank, device=target_device, dtype=torch.float32)
        )
        self.query_vh = nn.Parameter(
            torch.empty(self.rank, self.hidden_size, device=target_device, dtype=target_dtype)
        )
        self.query_bias = nn.Parameter(
            torch.zeros(self.q_head_size, device=target_device, dtype=target_dtype)
        )

        self.key_u = nn.Parameter(
            torch.empty(self.kv_head_size, self.rank, device=target_device, dtype=target_dtype)
        )
        self.key_s = nn.Parameter(torch.empty(self.rank, device=target_device, dtype=torch.float32))
        self.key_vh = nn.Parameter(
            torch.empty(self.rank, self.hidden_size, device=target_device, dtype=target_dtype)
        )
        self.key_bias = nn.Parameter(
            torch.zeros(self.kv_head_size, device=target_device, dtype=target_dtype)
        )

        self.value_u = nn.Parameter(
            torch.empty(self.kv_head_size, self.rank, device=target_device, dtype=target_dtype)
        )
        self.value_s = nn.Parameter(
            torch.empty(self.rank, device=target_device, dtype=torch.float32)
        )
        self.value_vh = nn.Parameter(
            torch.empty(self.rank, self.hidden_size, device=target_device, dtype=target_dtype)
        )
        self.value_bias = nn.Parameter(
            torch.zeros(self.kv_head_size, device=target_device, dtype=target_dtype)
        )

        self.out_proj_u = nn.Parameter(
            torch.empty(self.hidden_size, self.rank, device=target_device, dtype=target_dtype)
        )
        self.out_proj_s = nn.Parameter(
            torch.empty(self.rank, device=target_device, dtype=torch.float32)
        )
        self.out_proj_vh = nn.Parameter(
            torch.empty(self.rank, self.q_head_size, device=target_device, dtype=target_dtype)
        )
        self.out_proj_bias = nn.Parameter(
            torch.zeros(self.hidden_size, device=target_device, dtype=target_dtype)
        )

        # 4. Trainable LoRA Adapters
        if self.lora_rank > 0:
            self.q_lora_A = nn.Parameter(
                torch.empty(
                    self.lora_rank,
                    self.hidden_size,
                    device=target_device,
                    dtype=target_dtype,
                )
            )
            self.q_lora_B = nn.Parameter(
                torch.empty(
                    self.q_head_size,
                    self.lora_rank,
                    device=target_device,
                    dtype=target_dtype,
                )
            )
            self.k_lora_A = nn.Parameter(
                torch.empty(
                    self.lora_rank,
                    self.hidden_size,
                    device=target_device,
                    dtype=target_dtype,
                )
            )
            self.k_lora_B = nn.Parameter(
                torch.empty(
                    self.kv_head_size,
                    self.lora_rank,
                    device=target_device,
                    dtype=target_dtype,
                )
            )
            self.v_lora_A = nn.Parameter(
                torch.empty(
                    self.lora_rank,
                    self.hidden_size,
                    device=target_device,
                    dtype=target_dtype,
                )
            )
            self.v_lora_B = nn.Parameter(
                torch.empty(
                    self.kv_head_size,
                    self.lora_rank,
                    device=target_device,
                    dtype=target_dtype,
                )
            )
            self.out_lora_A = nn.Parameter(
                torch.empty(
                    self.lora_rank,
                    self.q_head_size,
                    device=target_device,
                    dtype=target_dtype,
                )
            )
            self.out_lora_B = nn.Parameter(
                torch.empty(
                    self.hidden_size,
                    self.lora_rank,
                    device=target_device,
                    dtype=target_dtype,
                )
            )

            for p_a in (self.q_lora_A, self.k_lora_A, self.v_lora_A, self.out_lora_A):
                nn.init.kaiming_uniform_(p_a, a=math.sqrt(5))
            for p_b in (self.q_lora_B, self.k_lora_B, self.v_lora_B, self.out_lora_B):
                nn.init.zeros_(p_b)
        else:
            self.register_parameter("q_lora_A", None)
            self.register_parameter("q_lora_B", None)
            self.register_parameter("k_lora_A", None)
            self.register_parameter("k_lora_B", None)
            self.register_parameter("v_lora_A", None)
            self.register_parameter("v_lora_B", None)
            self.register_parameter("out_lora_A", None)
            self.register_parameter("out_lora_B", None)

        self._init_random()

        # Freeze base SVD parameters
        for p in (
            self.query_u,
            self.query_s,
            self.query_vh,
            self.query_bias,
            self.key_u,
            self.key_s,
            self.key_vh,
            self.key_bias,
            self.value_u,
            self.value_s,
            self.value_vh,
            self.value_bias,
            self.out_proj_u,
            self.out_proj_s,
            self.out_proj_vh,
            self.out_proj_bias,
        ):
            p.requires_grad = False

        if old_attn is not None:
            self.copy_weights(old_attn)

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(config, "attention_probs_dropout_prob", dropout),
            )
            if config
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)
        self.triton_kwargs = {}

        if target_device is not None or target_dtype is not None:
            self.to(device=target_device, dtype=target_dtype)

    def _init_random(self) -> None:
        """Initialize orthogonal factor projections with identity scales."""

        def _ortho(t: nn.Parameter) -> None:
            with torch.no_grad():
                if t.device.type == "cuda" and t.dtype in (
                    torch.float16,
                    torch.bfloat16,
                ):
                    f = torch.empty(t.shape, device=t.device, dtype=torch.float32)
                    nn.init.orthogonal_(f)
                    t.data.copy_(f.to(dtype=t.dtype))
                else:
                    nn.init.orthogonal_(t)

        _ortho(self.query_u)
        nn.init.ones_(self.query_s)
        _ortho(self.query_vh)
        nn.init.zeros_(self.query_bias)

        _ortho(self.key_u)
        nn.init.ones_(self.key_s)
        _ortho(self.key_vh)
        nn.init.zeros_(self.key_bias)

        _ortho(self.value_u)
        nn.init.ones_(self.value_s)
        _ortho(self.value_vh)
        nn.init.zeros_(self.value_bias)

        _ortho(self.out_proj_u)
        nn.init.ones_(self.out_proj_s)
        _ortho(self.out_proj_vh)
        nn.init.zeros_(self.out_proj_bias)

    def _decompose_and_set(
        self,
        u_param: nn.Parameter,
        s_param: nn.Parameter,
        vh_param: nn.Parameter,
        bias_param: nn.Parameter,
        weight: torch.Tensor,
        bias: torch.Tensor | None,
    ) -> None:
        with torch.no_grad():
            W = weight.float()
            if W.shape != (u_param.shape[0], vh_param.shape[1]):
                logger.warning(
                    f"SVDLLMV2 shape mismatch: W={W.shape} vs expected=({u_param.shape[0]}, {vh_param.shape[1]}). Padding/truncating."
                )
                padded_W = torch.zeros(
                    u_param.shape[0], vh_param.shape[1], device=W.device, dtype=W.dtype
                )
                padded_W[
                    : min(u_param.shape[0], W.shape[0]),
                    : min(vh_param.shape[1], W.shape[1]),
                ] = W[
                    : min(u_param.shape[0], W.shape[0]),
                    : min(vh_param.shape[1], W.shape[1]),
                ]
                W = padded_W

            U, S, Vt = torch.linalg.svd(W.float(), full_matrices=False)
            total_energy = (S**2).sum().clamp(min=1e-12)
            cum_energy = torch.cumsum(S**2, dim=0) / total_energy
            r_energy = int((cum_energy < self.energy_threshold).sum().item()) + 1
            r = max(1, min(r_energy, self.rank, len(S)))

            u_param.data.zero_()
            s_param.data.zero_()
            vh_param.data.zero_()
            u_param.data[:, :r] = U[:, :r].to(dtype=u_param.dtype, device=u_param.device)
            s_param.data[:r] = S[:r].to(dtype=s_param.dtype, device=s_param.device)
            vh_param.data[:r, :] = Vt[:r, :].to(dtype=vh_param.dtype, device=vh_param.device)
            if bias is not None:
                bias_param.data.copy_(bias.to(dtype=bias_param.dtype, device=bias_param.device))

    def copy_weights(self, old_attn: nn.Module) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            src = old_attn
            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o, b_o = _unpack_weight_bias(old_attn.output.dense)

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
                if c_attn_w is not None:
                    split_size = c_attn_w.size(0) // 3
                    w_q = c_attn_w[:split_size, :]
                    w_k = c_attn_w[split_size : 2 * split_size, :]
                    w_v = c_attn_w[2 * split_size :, :]
                    if c_attn_b is not None:
                        b_q = c_attn_b[:split_size]
                        b_k = c_attn_b[split_size : 2 * split_size]
                        b_v = c_attn_b[2 * split_size :]

                w_o, b_o = _unpack_weight_bias(src.c_proj)

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q, b_q = _unpack_weight_bias(src.query)
                w_k, b_k = _unpack_weight_bias(src.key)
                w_v, b_v = _unpack_weight_bias(src.value)

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q, b_q = _unpack_weight_bias(src.q_proj)
                w_k, b_k = _unpack_weight_bias(src.k_proj)
                w_v, b_v = _unpack_weight_bias(src.v_proj)
                if hasattr(src, "o_proj"):
                    w_o, b_o = _unpack_weight_bias(src.o_proj)

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q, b_q = _unpack_weight_bias(src.q_lin)
                w_k, b_k = _unpack_weight_bias(src.k_lin)
                w_v, b_v = _unpack_weight_bias(src.v_lin)
                if hasattr(src, "out_lin"):
                    w_o, b_o = _unpack_weight_bias(src.out_lin)

            # Fallback: UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter(old_attn)
                    extracted = adapter.extract_qkv_out()
                    if extracted:
                        w_q = extracted.get("q_weight")
                        b_q = extracted.get("q_bias")
                        w_k = extracted.get("k_weight")
                        b_k = extracted.get("k_bias")
                        w_v = extracted.get("v_weight")
                        b_v = extracted.get("v_bias")
                        if w_o is None:
                            w_o = extracted.get("out_weight")
                            b_o = extracted.get("out_bias")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # Verification of essential projections
            missing = []
            if w_q is None:
                missing.append("query/q_proj")
            if w_k is None:
                missing.append("key/k_proj")
            if w_v is None:
                missing.append("value/v_proj")
            if missing:
                logger.warning(f"Failed to extract required projections: {missing}")
                return

            # Frobenius norm checks
            for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    logger.warning(f"Copied weight {name} has near-zero Frobenius norm (< 1e-4)")

            # Decompose and set Q
            self._decompose_and_set(
                self.query_u, self.query_s, self.query_vh, self.query_bias, w_q, b_q
            )

            # Decompose and set K & V with GQA pooling if needed
            if w_k.shape[0] == self.kv_head_size:
                self._decompose_and_set(
                    self.key_u, self.key_s, self.key_vh, self.key_bias, w_k, b_k
                )
            else:
                rep = w_k.shape[0] // self.kv_head_size
                w_k_pooled = (
                    w_k.view(self.num_kv_heads, rep, self.head_dim, -1)
                    .mean(dim=1)
                    .view(self.kv_head_size, -1)
                )
                b_k_pooled = (
                    b_k.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                    if b_k is not None
                    else None
                )
                self._decompose_and_set(
                    self.key_u,
                    self.key_s,
                    self.key_vh,
                    self.key_bias,
                    w_k_pooled,
                    b_k_pooled,
                )

            if w_v.shape[0] == self.kv_head_size:
                self._decompose_and_set(
                    self.value_u, self.value_s, self.value_vh, self.value_bias, w_v, b_v
                )
            else:
                rep = w_v.shape[0] // self.kv_head_size
                w_v_pooled = (
                    w_v.view(self.num_kv_heads, rep, self.head_dim, -1)
                    .mean(dim=1)
                    .view(self.kv_head_size, -1)
                )
                b_v_pooled = (
                    b_v.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                    if b_v is not None
                    else None
                )
                self._decompose_and_set(
                    self.value_u,
                    self.value_s,
                    self.value_vh,
                    self.value_bias,
                    w_v_pooled,
                    b_v_pooled,
                )

            # Decompose and set Out projection (if not nested encoder)
            if w_o is not None and not self.is_nested_encoder:
                if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                    logger.warning("Copied out_proj weight has near-zero Frobenius norm (< 1e-4)")
                self._decompose_and_set(
                    self.out_proj_u,
                    self.out_proj_s,
                    self.out_proj_vh,
                    self.out_proj_bias,
                    w_o,
                    b_o,
                )

    def _project(
        self,
        x: torch.Tensor,
        vh: nn.Parameter,
        s: nn.Parameter,
        u: nn.Parameter,
        b: nn.Parameter,
        lora_A: nn.Parameter | None,
        lora_B: nn.Parameter | None,
    ) -> torch.Tensor:
        """Factorized SVD projection with LoRA residual addition."""
        V_eff = s.to(dtype=vh.dtype).view(-1, 1) * vh
        base_out = triton_dispatch(self.triton_kernel, x, u, V_eff, bias=b)
        if self.lora_rank > 0 and lora_A is not None and lora_B is not None:
            lora_out = F.linear(F.linear(x, lora_A), lora_B) * self.scaling
            base_out = base_out + lora_out
        return base_out

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, seq_len, _ = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        # Defensive device check
        if self.query_u.device != device or self.query_u.dtype != dtype:
            self.to(device=device, dtype=dtype)

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Factorized linear projections with LoRA
        Q_act = self._project(
            hidden_states,
            self.query_vh,
            self.query_s,
            self.query_u,
            self.query_bias,
            self.q_lora_A,
            self.q_lora_B,
        )
        K_act = self._project(
            kv_input,
            self.key_vh,
            self.key_s,
            self.key_u,
            self.key_bias,
            self.k_lora_A,
            self.k_lora_B,
        )
        V_act = self._project(
            kv_input,
            self.value_vh,
            self.value_s,
            self.value_u,
            self.value_bias,
            self.v_lora_A,
            self.v_lora_B,
        )

        Q = (
            Q_act.view(batch_size, seq_len, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            K_act.view(batch_size, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V = (
            V_act.view(batch_size, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. RoPE support
        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
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # 3. GQA expansion
        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        # 4. Attention scores in FP32
        scale = 1.0 / math.sqrt(self.head_dim)
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_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 = getattr(self, "is_decoder", False) and not is_cross and seq_len > 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(batch_size, seq_len, self.all_head_size)
                )
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                if self.is_nested_encoder:
                    out = _sdpa_out
                else:
                    out = self._project(
                        _sdpa_out,
                        self.out_proj_vh,
                        self.out_proj_s,
                        self.out_proj_u,
                        self.out_proj_bias,
                        self.out_lora_A,
                        self.out_lora_B,
                    )
                if out.shape[-1] != hidden_states.shape[-1]:
                    if out.shape[-1] < hidden_states.shape[-1]:
                        out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
                    else:
                        out = out[..., : hidden_states.shape[-1]]
                return out, None
            except Exception:
                pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        # 5. Masking
        if self.is_decoder and not is_cross and seq_len > 1:
            causal_mask = (
                torch.triu(
                    torch.full(
                        (seq_len, kv_len),
                        float("-inf"),
                        device=scores.device,
                        dtype=torch.float32,
                    ),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask.unsqueeze(1).unsqueeze(2), -10000.0)
                elif torch.is_floating_point(eff_mask) and (eff_mask < 0).any():
                    scores = scores + eff_mask.unsqueeze(1).unsqueeze(2).float()
                else:
                    scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask.unsqueeze(1), -10000.0)
                else:
                    scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                if eff_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~eff_mask, -10000.0)
                else:
                    scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        probs = self.dropout(probs)

        # 6. Safe head_mask application
        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            probs = probs * head_mask

        # 7. Context computation & Head Merge
        context = (
            torch.matmul(probs, V)
            .transpose(1, 2)
            .contiguous()
            .view(batch_size, seq_len, self.all_head_size)
        )
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # 8. Output projection (bypassed if nested encoder)
        if self.is_nested_encoder:
            out = context
        else:
            out = self._project(
                context,
                self.out_proj_vh,
                self.out_proj_s,
                self.out_proj_u,
                self.out_proj_bias,
                self.out_lora_A,
                self.out_lora_B,
            )

        if out.shape[-1] != hidden_states.shape[-1]:
            if out.shape[-1] < hidden_states.shape[-1]:
                out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
            else:
                out = out[..., : hidden_states.shape[-1]]

        return out, (probs if output_attentions else None)

SVD-factorized attention layer with explicit U, Sigma, V^H matrix buffers and trainable LoRA recovery adapters.

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_gqa : bool

The type of the None singleton.

var triton_kernel : str

The type of the None singleton.

Methods

def copy_weights(self, old_attn: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        src = old_attn
        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o, b_o = _unpack_weight_bias(old_attn.output.dense)

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
            if c_attn_w is not None:
                split_size = c_attn_w.size(0) // 3
                w_q = c_attn_w[:split_size, :]
                w_k = c_attn_w[split_size : 2 * split_size, :]
                w_v = c_attn_w[2 * split_size :, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:split_size]
                    b_k = c_attn_b[split_size : 2 * split_size]
                    b_v = c_attn_b[2 * split_size :]

            w_o, b_o = _unpack_weight_bias(src.c_proj)

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q, b_q = _unpack_weight_bias(src.query)
            w_k, b_k = _unpack_weight_bias(src.key)
            w_v, b_v = _unpack_weight_bias(src.value)

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q, b_q = _unpack_weight_bias(src.q_proj)
            w_k, b_k = _unpack_weight_bias(src.k_proj)
            w_v, b_v = _unpack_weight_bias(src.v_proj)
            if hasattr(src, "o_proj"):
                w_o, b_o = _unpack_weight_bias(src.o_proj)

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q, b_q = _unpack_weight_bias(src.q_lin)
            w_k, b_k = _unpack_weight_bias(src.k_lin)
            w_v, b_v = _unpack_weight_bias(src.v_lin)
            if hasattr(src, "out_lin"):
                w_o, b_o = _unpack_weight_bias(src.out_lin)

        # Fallback: UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter(old_attn)
                extracted = adapter.extract_qkv_out()
                if extracted:
                    w_q = extracted.get("q_weight")
                    b_q = extracted.get("q_bias")
                    w_k = extracted.get("k_weight")
                    b_k = extracted.get("k_bias")
                    w_v = extracted.get("v_weight")
                    b_v = extracted.get("v_bias")
                    if w_o is None:
                        w_o = extracted.get("out_weight")
                        b_o = extracted.get("out_bias")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # Verification of essential projections
        missing = []
        if w_q is None:
            missing.append("query/q_proj")
        if w_k is None:
            missing.append("key/k_proj")
        if w_v is None:
            missing.append("value/v_proj")
        if missing:
            logger.warning(f"Failed to extract required projections: {missing}")
            return

        # Frobenius norm checks
        for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                logger.warning(f"Copied weight {name} has near-zero Frobenius norm (< 1e-4)")

        # Decompose and set Q
        self._decompose_and_set(
            self.query_u, self.query_s, self.query_vh, self.query_bias, w_q, b_q
        )

        # Decompose and set K & V with GQA pooling if needed
        if w_k.shape[0] == self.kv_head_size:
            self._decompose_and_set(
                self.key_u, self.key_s, self.key_vh, self.key_bias, w_k, b_k
            )
        else:
            rep = w_k.shape[0] // self.kv_head_size
            w_k_pooled = (
                w_k.view(self.num_kv_heads, rep, self.head_dim, -1)
                .mean(dim=1)
                .view(self.kv_head_size, -1)
            )
            b_k_pooled = (
                b_k.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                if b_k is not None
                else None
            )
            self._decompose_and_set(
                self.key_u,
                self.key_s,
                self.key_vh,
                self.key_bias,
                w_k_pooled,
                b_k_pooled,
            )

        if w_v.shape[0] == self.kv_head_size:
            self._decompose_and_set(
                self.value_u, self.value_s, self.value_vh, self.value_bias, w_v, b_v
            )
        else:
            rep = w_v.shape[0] // self.kv_head_size
            w_v_pooled = (
                w_v.view(self.num_kv_heads, rep, self.head_dim, -1)
                .mean(dim=1)
                .view(self.kv_head_size, -1)
            )
            b_v_pooled = (
                b_v.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                if b_v is not None
                else None
            )
            self._decompose_and_set(
                self.value_u,
                self.value_s,
                self.value_vh,
                self.value_bias,
                w_v_pooled,
                b_v_pooled,
            )

        # Decompose and set Out projection (if not nested encoder)
        if w_o is not None and not self.is_nested_encoder:
            if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                logger.warning("Copied out_proj weight has near-zero Frobenius norm (< 1e-4)")
            self._decompose_and_set(
                self.out_proj_u,
                self.out_proj_s,
                self.out_proj_vh,
                self.out_proj_bias,
                w_o,
                b_o,
            )

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class SVDLinear (in_features: int,
out_features: int,
rank: int,
bias: bool = True,
freeze_svd: bool = False,
lora_rank: int | None = 16,
lora_alpha: float = 32.0,
use_fused_kernel: bool = False,
use_rms_norm: bool = False,
eps: float = 1e-06,
dtype: torch.dtype = torch.float32,
device: torch.device | None = None)

Expand source code

class SVDLinear(nn.Module):
    """
    Low-rank factorized linear layer:
        y = (x V^T) U^T + b + (alpha / r_lora) (x A^T) B^T
    """

    # SVDLA-FIX-12: this is NOT an attention module
    is_custom_attention: bool = False

    triton_kernel: str = "low_rank"

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        freeze_svd: bool = False,
        lora_rank: int | None = 16,
        lora_alpha: float = 32.0,
        use_fused_kernel: bool = False,
        use_rms_norm: bool = False,
        eps: float = 1e-6,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.rank = int(rank)
        self.freeze_svd = bool(freeze_svd)
        self.lora_rank = int(lora_rank) if lora_rank is not None else 16
        self.lora_alpha = float(lora_alpha)
        self.use_fused_kernel = bool(use_fused_kernel) and HAS_FUSED_KERNEL
        self.use_rms_norm = bool(use_rms_norm)
        self.eps = float(eps)

        self.V = nn.Parameter(torch.empty(self.rank, self.in_features, dtype=dtype, device=device))
        self.U = nn.Parameter(torch.empty(self.out_features, self.rank, dtype=dtype, device=device))

        if self.use_rms_norm:
            self.norm_weight = nn.Parameter(torch.ones(self.rank, dtype=dtype, device=device))
        else:
            self.register_parameter("norm_weight", None)

        if self.freeze_svd:
            self.U.requires_grad = False
            self.V.requires_grad = False

        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(
                torch.empty(self.lora_rank, self.in_features, dtype=dtype, device=device)
            )
            self.lora_B = nn.Parameter(
                torch.empty(self.out_features, self.lora_rank, dtype=dtype, device=device)
            )
            self.lora_scaling = self.lora_alpha / max(1, self.lora_rank)
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)
            self.lora_scaling = 1.0

        if bias:
            self.bias = nn.Parameter(torch.empty(self.out_features, dtype=dtype, device=device))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def get_weight(self) -> torch.Tensor:
        W = torch.matmul(self.U, self.V)
        if self.lora_A is not None and self.lora_B is not None:
            W = W + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
        return W

    def reset_parameters(self) -> None:
        for param in (self.U, self.V):
            orig_dtype = param.data.dtype
            if orig_dtype in (torch.float16, torch.bfloat16):
                param.data = param.data.float()
                nn.init.orthogonal_(param)
                param.data = param.data.to(orig_dtype)
            else:
                nn.init.orthogonal_(param)
        if self.norm_weight is not None:
            nn.init.ones_(self.norm_weight)
        if self.lora_A is not None and self.lora_B is not None:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def _safe_triton_dispatch(self, x: torch.Tensor) -> torch.Tensor:
        """
        SVDLA-FIX-03: wrap triton_dispatch with try/except. On any failure
        (including native-level crashes that Python can catch), fall back to
        plain F.linear. This prevents STATUS_ACCESS_VIOLATION at high ranks.
        """
        if _HAS_TRITON and triton_dispatch is not None:
            try:
                return triton_dispatch(self.triton_kernel, x, self.U, self.V, bias=self.bias)
            except Exception as e:
                logger.warning(
                    "SVDLinear: triton_dispatch('%s') failed (%s); falling back to F.linear.",
                    self.triton_kernel,
                    e,
                )
        # Fallback: plain F.linear
        Vx = F.linear(x, self.V)
        return F.linear(Vx, self.U, self.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.use_fused_kernel and HAS_FUSED_KERNEL:
            try:
                out = fused_svd_linear(
                    x=x,
                    V=self.V,
                    U=self.U,
                    norm_weight=self.norm_weight if self.use_rms_norm else None,
                    bias=self.bias,
                    eps=self.eps,
                    use_rms_norm=self.use_rms_norm,
                )
            except Exception as e:
                logger.warning("SVDLinear: fused_svd_linear failed (%s); fallback.", e)
                Vx = F.linear(x, self.V)
                out = F.linear(Vx, self.U, self.bias)
        elif not self.use_rms_norm:
            out = self._safe_triton_dispatch(x)
        else:
            Vx = F.linear(x, self.V)
            if self.norm_weight is not None:
                norm = torch.rsqrt(Vx.pow(2).mean(-1, keepdim=True) + self.eps)
                Vx = Vx * norm * self.norm_weight
            out = F.linear(Vx, self.U, self.bias)

        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
            out = out + lora_out

        return out

    def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
        # SVDLA-FIX-15
        with torch.no_grad():
            r = min(self.rank, U_factor.size(1), V_factor.size(0))
            if r < self.rank:
                logger.debug(
                    "SVDLinear: factor rank %d < allocated rank %d; zeroing remainder.",
                    r,
                    self.rank,
                )
            out_d = min(self.out_features, U_factor.size(0))
            in_d = min(self.in_features, V_factor.size(1))
            self.U.data.zero_()
            self.V.data.zero_()
            self.U.data[:out_d, :r].copy_(U_factor[:out_d, :r].to(self.U.device, self.U.dtype))
            self.V.data[:r, :in_d].copy_(V_factor[:r, :in_d].to(self.V.device, self.V.dtype))

    def get_orthogonality_loss(self) -> torch.Tensor:
        U_orth = torch.matmul(self.U.t().contiguous(), self.U) - torch.eye(
            self.rank, device=self.U.device, dtype=self.U.dtype
        )
        V_orth = torch.matmul(self.V, self.V.t().contiguous()) - torch.eye(
            self.rank, device=self.V.device, dtype=self.V.dtype
        )
        return torch.norm(U_orth, p="fro") + torch.norm(V_orth, p="fro")

Low-rank factorized linear layer: y = (x V^T) U^T + b + (alpha / r_lora) (x A^T) B^T

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

var triton_kernel : str

The type of the None singleton.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self.use_fused_kernel and HAS_FUSED_KERNEL:
        try:
            out = fused_svd_linear(
                x=x,
                V=self.V,
                U=self.U,
                norm_weight=self.norm_weight if self.use_rms_norm else None,
                bias=self.bias,
                eps=self.eps,
                use_rms_norm=self.use_rms_norm,
            )
        except Exception as e:
            logger.warning("SVDLinear: fused_svd_linear failed (%s); fallback.", e)
            Vx = F.linear(x, self.V)
            out = F.linear(Vx, self.U, self.bias)
    elif not self.use_rms_norm:
        out = self._safe_triton_dispatch(x)
    else:
        Vx = F.linear(x, self.V)
        if self.norm_weight is not None:
            norm = torch.rsqrt(Vx.pow(2).mean(-1, keepdim=True) + self.eps)
            Vx = Vx * norm * self.norm_weight
        out = F.linear(Vx, self.U, self.bias)

    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
        out = out + lora_out

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def get_orthogonality_loss(self) ‑> torch.Tensor

Expand source code

def get_orthogonality_loss(self) -> torch.Tensor:
    U_orth = torch.matmul(self.U.t().contiguous(), self.U) - torch.eye(
        self.rank, device=self.U.device, dtype=self.U.dtype
    )
    V_orth = torch.matmul(self.V, self.V.t().contiguous()) - torch.eye(
        self.rank, device=self.V.device, dtype=self.V.dtype
    )
    return torch.norm(U_orth, p="fro") + torch.norm(V_orth, p="fro")

def get_weight(self) ‑> torch.Tensor

Expand source code

def get_weight(self) -> torch.Tensor:
    W = torch.matmul(self.U, self.V)
    if self.lora_A is not None and self.lora_B is not None:
        W = W + torch.matmul(self.lora_B, self.lora_A) * self.lora_scaling
    return W

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    for param in (self.U, self.V):
        orig_dtype = param.data.dtype
        if orig_dtype in (torch.float16, torch.bfloat16):
            param.data = param.data.float()
            nn.init.orthogonal_(param)
            param.data = param.data.to(orig_dtype)
        else:
            nn.init.orthogonal_(param)
    if self.norm_weight is not None:
        nn.init.ones_(self.norm_weight)
    if self.lora_A is not None and self.lora_B is not None:
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)
    if self.bias is not None:
        nn.init.zeros_(self.bias)

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None

Expand source code

def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None:
    # SVDLA-FIX-15
    with torch.no_grad():
        r = min(self.rank, U_factor.size(1), V_factor.size(0))
        if r < self.rank:
            logger.debug(
                "SVDLinear: factor rank %d < allocated rank %d; zeroing remainder.",
                r,
                self.rank,
            )
        out_d = min(self.out_features, U_factor.size(0))
        in_d = min(self.in_features, V_factor.size(1))
        self.U.data.zero_()
        self.V.data.zero_()
        self.U.data[:out_d, :r].copy_(U_factor[:out_d, :r].to(self.U.device, self.U.dtype))
        self.V.data[:r, :in_d].copy_(V_factor[:r, :in_d].to(self.V.device, self.V.dtype))

class SVDLinearAttention (config: Any | None = None,
old_attn: nn.Module | None = None,
rank: int | None = None,
mode: str = 'exact',
freeze_svd: bool = False,
lora_rank: int | None = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
use_decay_gating: bool = False,
decay_activation: str = 'softplus',
decay_channel_wise: bool = False,
use_fused_kernel: bool = False,
use_rms_norm: bool = False,
chunk_size: int = 64,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "svd_linear_attn",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "svd_linear_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
@ATTENTION_REGISTRY.register(
    "svd_linear", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
class SVDLinearAttention(BaseAttention):
    """
    SVD Factorized Linear Attention with optional exponential decay gating.
    """

    supports_gqa: bool = True
    supports_causal: bool = True
    supports_bidirectional: bool = True
    is_custom_attention: bool = True

    triton_kernel: str = "linear_attn"

    _VALID_MODES = {
        "exact",
        "linear",
        "decay_linear",
        "linear_decay",
        "decay",
        "gated_linear",
    }

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = None,
        mode: str = "exact",
        freeze_svd: bool = False,
        lora_rank: int | None = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        use_decay_gating: bool = False,
        decay_activation: str = "softplus",
        decay_channel_wise: bool = False,
        use_fused_kernel: bool = False,
        use_rms_norm: bool = False,
        chunk_size: int = 64,
        **kwargs: Any,
    ) -> None:
        actual_rank = rank or _cfg_get(config, "svd_rank", _cfg_get(config, "rank", 192))
        super().__init__(config, rank=actual_rank)

        # ---- dimension detection ----
        detected_hidden = None
        detected_heads = None
        detected_head_dim = None
        if old_attn is not None:
            for attr in ("all_head_size", "embed_dim", "hidden_size", "d_model"):
                if hasattr(old_attn, attr):
                    detected_hidden = getattr(old_attn, attr)
                    break
            if detected_hidden is None:
                for proj_name in ("query", "q_proj", "q_lin"):
                    if hasattr(old_attn, proj_name):
                        proj = getattr(old_attn, proj_name)
                        if hasattr(proj, "in_features"):
                            detected_hidden = proj.in_features
                            break
                        if hasattr(proj, "weight"):
                            detected_hidden = proj.weight.shape[-1]
                            break
            for attr in ("num_attention_heads", "num_heads", "n_head", "n_heads"):
                if hasattr(old_attn, attr):
                    detected_heads = getattr(old_attn, attr)
                    break
            for attr in ("attention_head_size", "head_dim", "d_head"):
                if hasattr(old_attn, attr):
                    detected_head_dim = getattr(old_attn, attr)
                    break

        self.hidden_size = int(
            _cfg_get(
                config,
                "hidden_size",
                kwargs.get("hidden_size", kwargs.get("embed_dim", detected_hidden or 768)),
            )
        )
        self.num_heads = int(
            _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(
                    config,
                    "num_heads",
                    kwargs.get(
                        "num_attention_heads",
                        kwargs.get("num_heads", detected_heads or 12),
                    ),
                ),
            )
        )
        self.head_dim = int(
            _cfg_get(
                config,
                "head_dim",
                kwargs.get(
                    "head_dim",
                    detected_head_dim or (self.hidden_size // max(1, self.num_heads)),
                ),
            )
        )

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        elif hasattr(old_attn, "num_key_value_heads"):
            self.num_kv_heads = int(old_attn.num_key_value_heads)
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = self.num_heads // self.num_kv_heads
        self.kv_head_size = self.num_kv_heads * self.head_dim
        self.q_head_size = self.num_heads * self.head_dim
        self.all_head_size = self.q_head_size

        self.rank = actual_rank

        # SVDLA-FIX-09: normalize mode
        mode_norm = str(mode).strip().lower()
        if mode_norm not in self._VALID_MODES:
            logger.warning("SVDLinearAttention: unknown mode '%s'; defaulting to 'exact'.", mode)
            mode_norm = "exact"
        self.mode = mode_norm

        self.dropout_rate = float(
            _cfg_get(
                config,
                "attention_probs_dropout_prob",
                _cfg_get(config, "dropout", 0.1),
            )
        )
        self.eps = float(_cfg_get(config, "linear_attention_eps", 1e-6))
        self.triton_kwargs = {"phi": "elu+1", "eps": self.eps}

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(_cfg_get(config, "model_type", "") or "").lower()
            self.is_decoder = (
                bool(_cfg_get(config, "is_decoder", False))
                or any(
                    _mt.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "falcon",
                        "bloom",
                        "opt",
                        "gemma",
                        "phi",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                or (old_attn is not None and hasattr(old_attn, "c_attn"))
            )

        self.is_nested_encoder = False
        self.lora_rank = int(lora_rank) if lora_rank is not None else 16
        self.lora_alpha = float(lora_alpha)

        self.use_decay_gating = bool(
            use_decay_gating
            or _cfg_get(config, "use_decay_gating", False)
            or (self.mode in ("decay_linear", "linear_decay", "decay", "gated_linear"))
        )
        self.decay_activation = str(_cfg_get(config, "decay_activation", decay_activation))
        self.decay_channel_wise = bool(_cfg_get(config, "decay_channel_wise", decay_channel_wise))
        self.use_fused_kernel = bool(
            use_fused_kernel or _cfg_get(config, "use_fused_kernel", False)
        )
        self.use_rms_norm = bool(use_rms_norm or _cfg_get(config, "use_rms_norm", False))
        # SVDLA-FIX-07
        self.chunk_size = max(8, int(_cfg_get(config, "chunk_size", chunk_size)))

        # ---- projections ----
        self.q_proj = SVDLinear(
            self.hidden_size,
            self.q_head_size,
            rank=self.rank,
            freeze_svd=freeze_svd,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            use_fused_kernel=self.use_fused_kernel,
            use_rms_norm=self.use_rms_norm,
            eps=self.eps,
        )
        self.k_proj = SVDLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            freeze_svd=freeze_svd,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            use_fused_kernel=self.use_fused_kernel,
            use_rms_norm=self.use_rms_norm,
            eps=self.eps,
        )
        self.v_proj = SVDLinear(
            self.hidden_size,
            self.kv_head_size,
            rank=self.rank,
            freeze_svd=freeze_svd,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            use_fused_kernel=self.use_fused_kernel,
            use_rms_norm=self.use_rms_norm,
            eps=self.eps,
        )
        self.out_proj = SVDLinear(
            self.q_head_size,
            self.hidden_size,
            rank=self.rank,
            freeze_svd=freeze_svd,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            use_fused_kernel=self.use_fused_kernel,
            use_rms_norm=self.use_rms_norm,
            eps=self.eps,
        )

        gate_out_dim = self.all_head_size if self.decay_channel_wise else self.num_heads
        self.decay_gate = nn.Linear(self.hidden_size, gate_out_dim, bias=True)
        self.dropout = nn.Dropout(self.dropout_rate)
        self._init_decay_gate()

        # storage accounting
        self._orig_storage_bytes = 0
        self._comp_storage_bytes = 0

        if old_attn is not None:
            self.copy_weights(old_attn)

    # ------------------------------------------------------------------
    def _init_decay_gate(self) -> None:
        # SVDLA-FIX-08
        with torch.no_grad():
            nn.init.kaiming_uniform_(self.decay_gate.weight, a=math.sqrt(5))
            log_min = math.log(1.0)
            log_max = math.log(1024.0)
            if self.decay_channel_wise:
                decay_times = torch.exp(torch.linspace(log_min, log_max, self.num_heads))
                decay_times = decay_times.repeat_interleave(self.head_dim)
            else:
                decay_times = torch.exp(torch.linspace(log_min, log_max, self.num_heads))
            bias_target = -torch.log(decay_times)
            bias_target = torch.clamp(bias_target, min=-20.0, max=20.0)
            if bias_target.numel() == self.decay_gate.bias.numel():
                self.decay_gate.bias.copy_(bias_target)
            else:
                logger.debug(
                    "SVDLinearAttention: decay_gate bias size mismatch (%d vs %d); zeroing.",
                    bias_target.numel(),
                    self.decay_gate.bias.numel(),
                )
                nn.init.zeros_(self.decay_gate.bias)

    # ------------------------------------------------------------------
    def _factorize_and_set(
        self,
        target_linear: SVDLinear,
        weight: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> None:
        with torch.no_grad():
            dev = weight.device
            try:
                U, S, Vt = torch.linalg.svd(weight.detach().float(), full_matrices=False)
            except Exception:
                U, S, Vt = torch.linalg.svd(weight.detach().cpu().float(), full_matrices=False)
                U = U.to(dev)
                S = S.to(dev)
                Vt = Vt.to(dev)

            r = min(target_linear.rank, U.size(1), Vt.size(0))
            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-5))
            U_r = (U[:, :r] * sqrt_S.unsqueeze(0)).to(device=dev, dtype=weight.dtype)
            V_r = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(device=dev, dtype=weight.dtype)
            target_linear.set_low_rank_factors(U_r, V_r)
            if bias is not None and target_linear.bias is not None:
                target_linear.bias.data.copy_(
                    bias.to(target_linear.bias.device, target_linear.bias.dtype)
                )

    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: nn.Module) -> None:
        """
        SVDLA-FIX-01: never raise. On missing projections, log a warning
        and fall back to the randomly-initialized state.
        """
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # nested encoder
            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                src = old_attn.self
                self.is_nested_encoder = True
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # PyTorch MultiheadAttention
            if hasattr(src, "in_proj_weight") and src.in_proj_weight is not None:
                in_proj_w = getattr(src, "in_proj_weight", None)
                in_proj_b = getattr(src, "in_proj_bias", None)
                if in_proj_w is not None:
                    w_q, w_k, w_v = in_proj_w.chunk(3, dim=0)
                if in_proj_b is not None:
                    b_q, b_k, b_v = in_proj_b.chunk(3, dim=0)
                if hasattr(src, "out_proj"):
                    w_o, b_o = _unpack_weight_bias(src.out_proj)

            # GPT-2 fused
            elif hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
                if c_attn_w is not None:
                    if (
                        c_attn_w.shape[0] == self.hidden_size
                        and c_attn_w.shape[1] != self.hidden_size
                    ):
                        c_attn_w = c_attn_w.t().contiguous()
                    q_dim = self.q_head_size
                    k_dim = self.kv_head_size
                    v_dim = self.kv_head_size
                    w_q = c_attn_w[:q_dim, :]
                    w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                    w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
                    if c_attn_b is not None:
                        b_q = c_attn_b[:q_dim]
                        b_k = c_attn_b[q_dim : q_dim + k_dim]
                        b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]
                out_w, out_b = _unpack_weight_bias(src.c_proj)
                if (
                    out_w is not None
                    and out_w.shape[0] == self.hidden_size
                    and out_w.shape[1] != self.hidden_size
                ):
                    out_w = out_w.t().contiguous()
                w_o, b_o = out_w, out_b

            # BERT / RoBERTa nested
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q, b_q = _unpack_weight_bias(src.query)
                w_k, b_k = _unpack_weight_bias(src.key)
                w_v, b_v = _unpack_weight_bias(src.value)
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                            break

            # Qwen / LLaMA
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q, b_q = _unpack_weight_bias(src.q_proj)
                w_k, b_k = _unpack_weight_bias(src.k_proj)
                w_v, b_v = _unpack_weight_bias(src.v_proj)
                if hasattr(src, "o_proj"):
                    w_o, b_o = _unpack_weight_bias(src.o_proj)

            # DistilBERT
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q, b_q = _unpack_weight_bias(src.q_lin)
                w_k, b_k = _unpack_weight_bias(src.k_lin)
                w_v, b_v = _unpack_weight_bias(src.v_lin)
                if hasattr(src, "out_lin"):
                    w_o, b_o = _unpack_weight_bias(src.out_lin)

            # UniversalAdapter fallback
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    if w_q is None:
                        w_q, b_q = adapter.get_weight_bias(old_attn, "q")
                    if w_k is None:
                        w_k, b_k = adapter.get_weight_bias(old_attn, "k")
                    if w_v is None:
                        w_v, b_v = adapter.get_weight_bias(old_attn, "v")
                    if w_o is None and not self.is_nested_encoder:
                        w_o, b_o = adapter.get_weight_bias(old_attn, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # ---- missing check: WARN, do not raise ----
            missing = []
            if w_q is None:
                missing.append("q")
            if w_k is None:
                missing.append("k")
            if w_v is None:
                missing.append("v")
            if missing:
                logger.warning(
                    "SVDLinearAttention: could not extract %s from %s; "
                    "keeping random init. This will likely reduce quality.",
                    missing,
                    type(old_attn).__name__,
                )
                return  # early return; projections keep their SVD init

            # ---- near-zero norm check: WARN only ----
            for name, w in (("w_q", w_q), ("w_k", w_k), ("w_v", w_v)):
                if w is None:
                    continue
                try:
                    n = torch.norm(w.float(), p="fro").item()
                except Exception:
                    n = float("nan")
                if not math.isfinite(n) or n < 1e-4:
                    logger.warning(
                        "SVDLinearAttention: %s has near-zero/NaN norm (%.2e); proceeding anyway.",
                        name,
                        n,
                    )

            # ---- factorize Q ----
            try:
                self._factorize_and_set(self.q_proj, w_q, b_q)
            except Exception as e:
                logger.warning("SVDLinearAttention: Q factorization failed (%s).", e)
                return

            # ---- factorize K and V with GQA adaptation ----
            def _adapt_to_kv_size(weight, bias, mode="mean"):
                w, b = weight, bias
                if w.shape[0] == self.kv_head_size:
                    return w, b
                if w.shape[0] > self.kv_head_size and w.shape[0] % self.kv_head_size == 0:
                    rep = w.shape[0] // self.kv_head_size
                    if mode == "slice":
                        w = w.contiguous().view(self.num_kv_heads, rep, -1)[:, 0, :].contiguous()
                        if b is not None:
                            b = (
                                b.contiguous()
                                .view(self.num_kv_heads, rep, -1)[:, 0, :]
                                .contiguous()
                                .view(-1)
                            )
                    elif mode == "first":
                        w = w[: self.kv_head_size, :].contiguous()
                        if b is not None:
                            b = b[: self.kv_head_size].contiguous()
                    else:  # mean
                        w = w.contiguous().view(self.num_kv_heads, rep, -1).mean(dim=1)
                        if b is not None:
                            b = b.contiguous().view(self.num_kv_heads, rep, -1).mean(dim=1).view(-1)
                elif w.shape[0] < self.kv_head_size and self.kv_head_size % w.shape[0] == 0:
                    rep = self.kv_head_size // w.shape[0]
                    w = w.repeat(rep, 1)
                    if b is not None:
                        b = b.repeat(rep)
                else:
                    padded_w = torch.zeros(
                        self.kv_head_size,
                        w.shape[1],
                        device=w.device,
                        dtype=w.dtype,
                    )
                    min_rows = min(self.kv_head_size, w.shape[0])
                    padded_w[:min_rows, :] = w[:min_rows, :]
                    w = padded_w
                    if b is not None:
                        padded_b = torch.zeros(
                            self.kv_head_size,
                            device=b.device,
                            dtype=b.dtype,
                        )
                        padded_b[:min_rows] = b[:min_rows]
                        b = padded_b
                return w, b

            try:
                w_k_a, b_k_a = _adapt_to_kv_size(w_k, b_k, mode="slice")
                self._factorize_and_set(self.k_proj, w_k_a, b_k_a)
            except Exception as e:
                logger.warning("SVDLinearAttention: K factorization failed (%s).", e)
                return

            try:
                w_v_a, b_v_a = _adapt_to_kv_size(w_v, b_v, mode="slice")
                self._factorize_and_set(self.v_proj, w_v_a, b_v_a)
            except Exception as e:
                logger.warning("SVDLinearAttention: V factorization failed (%s).", e)
                return

            # ---- output projection ----
            if w_o is not None and not self.is_nested_encoder:
                try:
                    n_o = torch.norm(w_o.float(), p="fro").item()
                except Exception:
                    n_o = float("nan")
                if not math.isfinite(n_o) or n_o < 1e-4:
                    logger.warning(
                        "SVDLinearAttention: out_proj has near-zero/NaN norm (%.2e); "
                        "proceeding anyway.",
                        n_o,
                    )
                try:
                    self._factorize_and_set(self.out_proj, w_o, b_o)
                except Exception as e:
                    logger.warning("SVDLinearAttention: out_proj factorization failed (%s).", e)

    # ------------------------------------------------------------------
    # Forward paths
    # ------------------------------------------------------------------
    def _forward_decay_gated_linear_kernel(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        hidden_states: torch.Tensor,
        mask: torch.Tensor | None = None,
        past_state: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        B, H, N, D_k = q.shape
        D_v = v.shape[-1]

        compute_dtype = torch.bfloat16 if q.dtype == torch.bfloat16 else torch.float32

        phi_q = F.elu(q.to(compute_dtype)) + 1.0
        phi_k = F.elu(k.to(compute_dtype)) + 1.0
        v_f = v.to(compute_dtype)

        mask_expanded = _normalize_mask(mask, B, H, N, compute_dtype, q.device)
        if mask_expanded is not None:
            # mask_expanded shape (B, 1, N, 1) or (B, H, N, 1)
            if mask_expanded.shape[1] == 1:
                phi_k = phi_k * mask_expanded
            else:
                phi_k = phi_k * mask_expanded

        # decay
        decay_logits = self.decay_gate(hidden_states)
        if self.decay_channel_wise:
            decay_logits = decay_logits.view(B, N, H, D_k).permute(0, 2, 1, 3)
            lambda_decay = (
                compute_exponential_decay(decay_logits, activation=self.decay_activation)
                .unsqueeze(-1)
                .to(compute_dtype)
            )
        else:
            decay_logits = decay_logits.permute(0, 2, 1).unsqueeze(-1).unsqueeze(-1)
            lambda_decay = compute_exponential_decay(
                decay_logits, activation=self.decay_activation
            ).to(compute_dtype)

        # padded positions do not decay
        if mask_expanded is not None:
            m_dec = mask_expanded
            while m_dec.dim() < lambda_decay.dim():
                m_dec = m_dec.unsqueeze(-1)
            lambda_decay = torch.where(m_dec > 0, lambda_decay, torch.ones_like(lambda_decay))

        # ---- chunked scan for non-channel-wise and no past state ----
        if past_state is None and not self.decay_channel_wise:
            chunk_size = self.chunk_size
            num_chunks = (N + chunk_size - 1) // chunk_size
            S_chunk = torch.zeros(B, H, D_k, D_v, device=q.device, dtype=compute_dtype)
            Z_chunk = torch.zeros(B, H, D_k, 1, device=q.device, dtype=compute_dtype)
            outputs = []

            for c in range(num_chunks):
                start = c * chunk_size
                end = min(start + chunk_size, N)
                qc = phi_q[:, :, start:end, :]
                kc = phi_k[:, :, start:end, :]
                vc = v_f[:, :, start:end, :]
                lam_c = lambda_decay[:, :, start:end].squeeze(-1).squeeze(-1)
                Lc = end - start

                # SVDLA-FIX-11: safe log with floor
                log_lam_c = torch.log(torch.clamp(lam_c, min=1e-6, max=1.0))
                log_gamma_c = torch.cumsum(log_lam_c, dim=-1)

                log_M_intra = log_gamma_c.unsqueeze(-1) - log_gamma_c.unsqueeze(-2)
                causal_bool = torch.tril(torch.ones(Lc, Lc, dtype=torch.bool, device=qc.device))
                M_intra = torch.exp(torch.clamp(log_M_intra, max=0.0)) * causal_bool.unsqueeze(
                    0
                ).unsqueeze(0)

                aff_intra = torch.matmul(qc, kc.transpose(-1, -2).contiguous()) * M_intra
                num_intra = torch.matmul(aff_intra, vc)
                den_intra = aff_intra.sum(dim=-1, keepdim=True)

                decay_from_past = torch.exp(torch.clamp(log_gamma_c, max=0.0)).unsqueeze(-1)
                num_inter = torch.matmul(qc, S_chunk) * decay_from_past
                den_inter = torch.matmul(qc, Z_chunk) * decay_from_past

                out_c = (num_intra + num_inter) / (den_intra + den_inter + self.eps)
                outputs.append(out_c)

                total_chunk_decay = torch.exp(
                    torch.clamp(log_gamma_c[:, :, -1:], max=0.0)
                ).unsqueeze(-1)
                log_w_to_end = log_gamma_c[:, :, -1:] - log_gamma_c
                w_to_end = torch.exp(torch.clamp(log_w_to_end, max=0.0)).unsqueeze(-1)
                kc_w = kc * w_to_end

                S_chunk = total_chunk_decay * S_chunk + torch.matmul(
                    kc_w.transpose(-1, -2).contiguous(), vc
                )
                Z_chunk = (
                    total_chunk_decay * Z_chunk
                    + kc_w.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous()
                )

            context = torch.cat(outputs, dim=2).to(q.dtype)
            return context, (S_chunk, Z_chunk)

        # ---- sequential recurrent scan ----
        if past_state is not None:
            S, Z = past_state
            S = S.to(device=q.device, dtype=compute_dtype)
            Z = Z.to(device=q.device, dtype=compute_dtype)
        else:
            S = torch.zeros(B, H, D_k, D_v, device=q.device, dtype=compute_dtype)
            Z = torch.zeros(B, H, D_k, 1, device=q.device, dtype=compute_dtype)

        outputs = []
        for t in range(N):
            qt = phi_q[:, :, t : t + 1, :]
            kt = phi_k[:, :, t : t + 1, :]
            vt = v_f[:, :, t : t + 1, :]
            lam_t = lambda_decay[:, :, t]

            kt_t = kt.transpose(-1, -2).contiguous()
            kv_t = torch.matmul(kt_t, vt)
            S = lam_t * S + kv_t
            Z = lam_t * Z + kt_t

            num_t = torch.matmul(qt, S)
            den_t = torch.matmul(qt, Z) + self.eps
            outputs.append(num_t / den_t)

        context = torch.cat(outputs, dim=2).to(q.dtype)
        return context, (S, Z)

    def _forward_linear_kernel(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: torch.Tensor | None,
    ) -> torch.Tensor:
        if mask is None and _HAS_TRITON and triton_dispatch is not None:
            try:
                return triton_dispatch(
                    self.triton_kernel,
                    q,
                    k,
                    v,
                    is_causal=self.is_decoder,
                    **self.triton_kwargs,
                )
            except Exception as e:
                logger.warning(
                    "SVDLinearAttention: triton linear_attn failed (%s); fallback.",
                    e,
                )

        B, H, N, D_k = q.shape
        D_v = v.shape[-1]

        phi_q = F.elu(q.float()) + 1.0
        phi_k = F.elu(k.float()) + 1.0
        v_f = v.float()

        mask_expanded = _normalize_mask(mask, B, H, N, torch.float32, q.device)
        if mask_expanded is not None:
            phi_k = phi_k * mask_expanded

        if self.is_decoder and N > 1:
            chunk_size = self.chunk_size
            num_chunks = (N + chunk_size - 1) // chunk_size
            S_chunk = torch.zeros(B, H, D_k, D_v, device=q.device, dtype=torch.float32)
            Z_chunk = torch.zeros(B, H, D_k, 1, device=q.device, dtype=torch.float32)
            outputs = []

            for c in range(num_chunks):
                start = c * chunk_size
                end = min(start + chunk_size, N)
                qc = phi_q[:, :, start:end, :]
                kc = phi_k[:, :, start:end, :]
                vc = v_f[:, :, start:end, :]
                Lc = end - start

                aff_intra = torch.matmul(qc, kc.transpose(-1, -2).contiguous())
                causal_intra = (
                    torch.tril(torch.ones(Lc, Lc, device=q.device, dtype=torch.bool))
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                aff_intra = aff_intra * causal_intra
                num_intra = torch.matmul(aff_intra, vc)
                den_intra = aff_intra.sum(dim=-1, keepdim=True)

                num_inter = torch.matmul(qc, S_chunk)
                den_inter = torch.matmul(qc, Z_chunk)

                out_c = (num_intra + num_inter) / (den_intra + den_inter + self.eps)
                outputs.append(out_c)

                S_chunk = S_chunk + torch.matmul(kc.transpose(-1, -2).contiguous(), vc)
                Z_chunk = Z_chunk + kc.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous()

            context = torch.cat(outputs, dim=2).to(q.dtype)
        else:
            kv_state = torch.matmul(phi_k.transpose(-1, -2).contiguous(), v_f)
            numerator = torch.matmul(phi_q, kv_state)
            k_sum = phi_k.sum(dim=-2, keepdim=True)
            denominator = torch.matmul(phi_q, k_sum.transpose(-1, -2).contiguous()) + self.eps
            context = (numerator / denominator).to(q.dtype)

        return context

    def _forward_exact_softmax(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: torch.Tensor | None,
        is_cross: bool = False,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        is_causal_attn = self.is_decoder and not is_cross and q.shape[-2] > 1

        if not output_attentions:
            try:
                sdpa_mask = None
                can_use_sdpa = False
                if mask is None:
                    can_use_sdpa = True
                elif mask.dim() == 2 and not is_causal_attn:
                    sdpa_mask = mask.bool().view(mask.shape[0], 1, 1, mask.shape[1])
                    can_use_sdpa = True
                elif mask.dim() == 4 and not is_causal_attn:
                    sdpa_mask = mask.bool() if mask.dtype == torch.bool else mask
                    can_use_sdpa = True

                if can_use_sdpa:
                    drop_p = self.dropout.p if self.training else 0.0
                    context = F.scaled_dot_product_attention(
                        q,
                        k,
                        v,
                        attn_mask=sdpa_mask,
                        dropout_p=drop_p,
                        is_causal=is_causal_attn if sdpa_mask is None else False,
                    )
                    return context, None
            except Exception:
                pass

        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(q.float(), k.float().transpose(-1, -2).contiguous()) * scale

        if is_causal_attn:
            causal_mask = (
                torch.triu(
                    torch.full(
                        (q.shape[-2], k.shape[-2]),
                        float("-inf"),
                        device=scores.device,
                        dtype=torch.float32,
                    ),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        if mask is not None:
            scores = apply_attention_mask(scores, mask)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(q.dtype)
        probs = self.dropout(probs)
        context = torch.matmul(probs, v)
        return context, probs

    # ------------------------------------------------------------------
    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,
        *args: Any,
        **kwargs: Any,
    ):
        query = kwargs.get("query")
        key = kwargs.get("key")
        mask = kwargs.get("mask")
        return_state = kwargs.get("return_state", False)
        past_state = kwargs.get("past_state")

        if return_state and (
            not self.use_decay_gating
            and self.mode not in ("decay_linear", "linear_decay", "decay", "gated_linear")
        ):
            raise ValueError(
                f"return_state=True is only supported in decay/gated modes, got mode='{self.mode}'"
            )

        if hidden_states is None:
            hidden_states = query if query is not None else key
        if attention_mask is None:
            attention_mask = mask

        B, N, _ = hidden_states.shape
        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        kv_len = kv_input.shape[1]

        # 1. projections
        q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_proj(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. RoPE
        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
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # 3. GQA repeat
        if self.num_kv_heads < self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        updated_state: tuple[torch.Tensor, torch.Tensor] | None = None

        # 4. dispatch
        if self.mode == "exact":
            context, attn_probs = self._forward_exact_softmax(
                q,
                k,
                v,
                eff_mask,
                is_cross=is_cross,
                output_attentions=bool(output_attentions),
            )
        elif self.use_decay_gating or self.mode in (
            "decay_linear",
            "linear_decay",
            "decay",
            "gated_linear",
        ):
            context, updated_state = self._forward_decay_gated_linear_kernel(
                q=q,
                k=k,
                v=v,
                hidden_states=hidden_states,
                mask=eff_mask,
                past_state=past_state,
            )
            attn_probs = None
        else:
            context = self._forward_linear_kernel(q, k, v, eff_mask)
            attn_probs = None

        # 5. head_mask
        if head_mask is not None:
            hm = head_mask
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            hm = hm.to(context.dtype)
            if attn_probs is not None:
                attn_probs = attn_probs * hm
            context = context * hm

        # 6. head merge & output projection
        context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)

        # SVDLA-FIX-06: explicit error on shape mismatch
        if output.shape[-1] != hidden_states.shape[-1]:
            raise RuntimeError(
                f"SVDLinearAttention: output dim {output.shape[-1]} != "
                f"input hidden dim {hidden_states.shape[-1]}"
            )

        attn_out = attn_probs if output_attentions else None

        if return_state and updated_state is not None:
            return output, attn_out, updated_state
        return output, attn_out

    # ------------------------------------------------------------------
    def step_recurrent(
        self,
        x_t: torch.Tensor,
        past_state: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        if x_t.dim() == 2:
            x_t = x_t.unsqueeze(1)

        B = x_t.shape[0]
        q = self.q_proj(x_t).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        k = (
            self.k_proj(x_t)
            .view(B, 1, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(x_t)
            .view(B, 1, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        if self.num_kv_heads < self.num_heads:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        phi_q = F.elu(q.float()) + 1.0
        phi_k = F.elu(k.float()) + 1.0

        _, H, _, D_k = q.shape
        D_v = v.shape[-1]

        if self.use_decay_gating:
            decay_logits = self.decay_gate(x_t)
            if self.decay_channel_wise:
                decay_logits = decay_logits.view(B, 1, H, D_k).permute(0, 2, 1, 3)
                lam_t = (
                    compute_exponential_decay(decay_logits, activation=self.decay_activation)
                    .unsqueeze(-1)
                    .squeeze(2)
                )
            else:
                decay_logits = decay_logits.permute(0, 2, 1).unsqueeze(-1).unsqueeze(-1)
                lam_t = compute_exponential_decay(
                    decay_logits, activation=self.decay_activation
                ).squeeze(2)
        else:
            lam_t = torch.ones(B, H, 1, 1, device=x_t.device, dtype=torch.float32)

        if past_state is not None:
            S_prev, Z_prev = past_state
        else:
            S_prev = torch.zeros(B, H, D_k, D_v, device=x_t.device, dtype=torch.float32)
            Z_prev = torch.zeros(B, H, D_k, 1, device=x_t.device, dtype=torch.float32)

        kt_t = phi_k.transpose(-1, -2).contiguous()
        kv_t = torch.matmul(kt_t, v.float())

        S_t = lam_t * S_prev + kv_t
        Z_t = lam_t * Z_prev + kt_t

        num_t = torch.matmul(phi_q, S_t)
        den_t = torch.matmul(phi_q, Z_t) + self.eps
        y_t = num_t / den_t

        context = y_t.transpose(1, 2).contiguous().view(B, 1, self.all_head_size).to(x_t.dtype)
        if self.is_nested_encoder:
            out_t = context
        else:
            out_t = self.out_proj(context)

        if out_t.shape[-1] != x_t.shape[-1]:
            raise RuntimeError(
                f"SVDLinearAttention.step_recurrent: output dim "
                f"{out_t.shape[-1]} != input dim {x_t.shape[-1]}"
            )

        return out_t, (S_t, Z_t)

    def get_orthogonality_loss(self) -> torch.Tensor:
        return (
            self.q_proj.get_orthogonality_loss()
            + self.k_proj.get_orthogonality_loss()
            + self.v_proj.get_orthogonality_loss()
            + self.out_proj.get_orthogonality_loss()
        )

SVD Factorized Linear Attention with optional exponential decay gating.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• BaseAttention
• torch.nn.modules.module.Module

Subclasses

• ProjectionFactorizedAttention
• DecayGatedLinearAttention

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 triton_kernel : str

The type of the None singleton.

Methods

def copy_weights(self, old_attn: nn.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    """
    SVDLA-FIX-01: never raise. On missing projections, log a warning
    and fall back to the randomly-initialized state.
    """
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # nested encoder
        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            src = old_attn.self
            self.is_nested_encoder = True
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight
                b_o = old_attn.output.dense.bias
        else:
            src = old_attn

        # PyTorch MultiheadAttention
        if hasattr(src, "in_proj_weight") and src.in_proj_weight is not None:
            in_proj_w = getattr(src, "in_proj_weight", None)
            in_proj_b = getattr(src, "in_proj_bias", None)
            if in_proj_w is not None:
                w_q, w_k, w_v = in_proj_w.chunk(3, dim=0)
            if in_proj_b is not None:
                b_q, b_k, b_v = in_proj_b.chunk(3, dim=0)
            if hasattr(src, "out_proj"):
                w_o, b_o = _unpack_weight_bias(src.out_proj)

        # GPT-2 fused
        elif hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
            if c_attn_w is not None:
                if (
                    c_attn_w.shape[0] == self.hidden_size
                    and c_attn_w.shape[1] != self.hidden_size
                ):
                    c_attn_w = c_attn_w.t().contiguous()
                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size
                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]
            out_w, out_b = _unpack_weight_bias(src.c_proj)
            if (
                out_w is not None
                and out_w.shape[0] == self.hidden_size
                and out_w.shape[1] != self.hidden_size
            ):
                out_w = out_w.t().contiguous()
            w_o, b_o = out_w, out_b

        # BERT / RoBERTa nested
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q, b_q = _unpack_weight_bias(src.query)
            w_k, b_k = _unpack_weight_bias(src.key)
            w_v, b_v = _unpack_weight_bias(src.value)
            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                        break

        # Qwen / LLaMA
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q, b_q = _unpack_weight_bias(src.q_proj)
            w_k, b_k = _unpack_weight_bias(src.k_proj)
            w_v, b_v = _unpack_weight_bias(src.v_proj)
            if hasattr(src, "o_proj"):
                w_o, b_o = _unpack_weight_bias(src.o_proj)

        # DistilBERT
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q, b_q = _unpack_weight_bias(src.q_lin)
            w_k, b_k = _unpack_weight_bias(src.k_lin)
            w_v, b_v = _unpack_weight_bias(src.v_lin)
            if hasattr(src, "out_lin"):
                w_o, b_o = _unpack_weight_bias(src.out_lin)

        # UniversalAdapter fallback
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                if w_q is None:
                    w_q, b_q = adapter.get_weight_bias(old_attn, "q")
                if w_k is None:
                    w_k, b_k = adapter.get_weight_bias(old_attn, "k")
                if w_v is None:
                    w_v, b_v = adapter.get_weight_bias(old_attn, "v")
                if w_o is None and not self.is_nested_encoder:
                    w_o, b_o = adapter.get_weight_bias(old_attn, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # ---- missing check: WARN, do not raise ----
        missing = []
        if w_q is None:
            missing.append("q")
        if w_k is None:
            missing.append("k")
        if w_v is None:
            missing.append("v")
        if missing:
            logger.warning(
                "SVDLinearAttention: could not extract %s from %s; "
                "keeping random init. This will likely reduce quality.",
                missing,
                type(old_attn).__name__,
            )
            return  # early return; projections keep their SVD init

        # ---- near-zero norm check: WARN only ----
        for name, w in (("w_q", w_q), ("w_k", w_k), ("w_v", w_v)):
            if w is None:
                continue
            try:
                n = torch.norm(w.float(), p="fro").item()
            except Exception:
                n = float("nan")
            if not math.isfinite(n) or n < 1e-4:
                logger.warning(
                    "SVDLinearAttention: %s has near-zero/NaN norm (%.2e); proceeding anyway.",
                    name,
                    n,
                )

        # ---- factorize Q ----
        try:
            self._factorize_and_set(self.q_proj, w_q, b_q)
        except Exception as e:
            logger.warning("SVDLinearAttention: Q factorization failed (%s).", e)
            return

        # ---- factorize K and V with GQA adaptation ----
        def _adapt_to_kv_size(weight, bias, mode="mean"):
            w, b = weight, bias
            if w.shape[0] == self.kv_head_size:
                return w, b
            if w.shape[0] > self.kv_head_size and w.shape[0] % self.kv_head_size == 0:
                rep = w.shape[0] // self.kv_head_size
                if mode == "slice":
                    w = w.contiguous().view(self.num_kv_heads, rep, -1)[:, 0, :].contiguous()
                    if b is not None:
                        b = (
                            b.contiguous()
                            .view(self.num_kv_heads, rep, -1)[:, 0, :]
                            .contiguous()
                            .view(-1)
                        )
                elif mode == "first":
                    w = w[: self.kv_head_size, :].contiguous()
                    if b is not None:
                        b = b[: self.kv_head_size].contiguous()
                else:  # mean
                    w = w.contiguous().view(self.num_kv_heads, rep, -1).mean(dim=1)
                    if b is not None:
                        b = b.contiguous().view(self.num_kv_heads, rep, -1).mean(dim=1).view(-1)
            elif w.shape[0] < self.kv_head_size and self.kv_head_size % w.shape[0] == 0:
                rep = self.kv_head_size // w.shape[0]
                w = w.repeat(rep, 1)
                if b is not None:
                    b = b.repeat(rep)
            else:
                padded_w = torch.zeros(
                    self.kv_head_size,
                    w.shape[1],
                    device=w.device,
                    dtype=w.dtype,
                )
                min_rows = min(self.kv_head_size, w.shape[0])
                padded_w[:min_rows, :] = w[:min_rows, :]
                w = padded_w
                if b is not None:
                    padded_b = torch.zeros(
                        self.kv_head_size,
                        device=b.device,
                        dtype=b.dtype,
                    )
                    padded_b[:min_rows] = b[:min_rows]
                    b = padded_b
            return w, b

        try:
            w_k_a, b_k_a = _adapt_to_kv_size(w_k, b_k, mode="slice")
            self._factorize_and_set(self.k_proj, w_k_a, b_k_a)
        except Exception as e:
            logger.warning("SVDLinearAttention: K factorization failed (%s).", e)
            return

        try:
            w_v_a, b_v_a = _adapt_to_kv_size(w_v, b_v, mode="slice")
            self._factorize_and_set(self.v_proj, w_v_a, b_v_a)
        except Exception as e:
            logger.warning("SVDLinearAttention: V factorization failed (%s).", e)
            return

        # ---- output projection ----
        if w_o is not None and not self.is_nested_encoder:
            try:
                n_o = torch.norm(w_o.float(), p="fro").item()
            except Exception:
                n_o = float("nan")
            if not math.isfinite(n_o) or n_o < 1e-4:
                logger.warning(
                    "SVDLinearAttention: out_proj has near-zero/NaN norm (%.2e); "
                    "proceeding anyway.",
                    n_o,
                )
            try:
                self._factorize_and_set(self.out_proj, w_o, b_o)
            except Exception as e:
                logger.warning("SVDLinearAttention: out_proj factorization failed (%s).", e)

SVDLA-FIX-01: never raise. On missing projections, log a warning and fall back to the randomly-initialized state.

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,
*args: Any,
**kwargs: Any) ‑> Callable[..., typing.Any]

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,
    *args: Any,
    **kwargs: Any,
):
    query = kwargs.get("query")
    key = kwargs.get("key")
    mask = kwargs.get("mask")
    return_state = kwargs.get("return_state", False)
    past_state = kwargs.get("past_state")

    if return_state and (
        not self.use_decay_gating
        and self.mode not in ("decay_linear", "linear_decay", "decay", "gated_linear")
    ):
        raise ValueError(
            f"return_state=True is only supported in decay/gated modes, got mode='{self.mode}'"
        )

    if hidden_states is None:
        hidden_states = query if query is not None else key
    if attention_mask is None:
        attention_mask = mask

    B, N, _ = hidden_states.shape
    is_cross = encoder_hidden_states is not None
    kv_input = encoder_hidden_states if is_cross else hidden_states
    kv_len = kv_input.shape[1]

    # 1. projections
    q = (
        self.q_proj(hidden_states)
        .view(B, N, self.num_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    k = (
        self.k_proj(kv_input)
        .view(B, kv_len, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v = (
        self.v_proj(kv_input)
        .view(B, kv_len, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # 2. RoPE
    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
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    # 3. GQA repeat
    if self.num_kv_heads < self.num_heads:
        k = repeat_kv(k, self.kv_repeat)
        v = repeat_kv(v, self.kv_repeat)

    eff_mask = encoder_attention_mask if is_cross else attention_mask
    updated_state: tuple[torch.Tensor, torch.Tensor] | None = None

    # 4. dispatch
    if self.mode == "exact":
        context, attn_probs = self._forward_exact_softmax(
            q,
            k,
            v,
            eff_mask,
            is_cross=is_cross,
            output_attentions=bool(output_attentions),
        )
    elif self.use_decay_gating or self.mode in (
        "decay_linear",
        "linear_decay",
        "decay",
        "gated_linear",
    ):
        context, updated_state = self._forward_decay_gated_linear_kernel(
            q=q,
            k=k,
            v=v,
            hidden_states=hidden_states,
            mask=eff_mask,
            past_state=past_state,
        )
        attn_probs = None
    else:
        context = self._forward_linear_kernel(q, k, v, eff_mask)
        attn_probs = None

    # 5. head_mask
    if head_mask is not None:
        hm = head_mask
        if hm.dim() == 1:
            hm = hm.view(1, -1, 1, 1)
        elif hm.dim() == 2:
            hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
        elif hm.dim() == 3:
            hm = hm.unsqueeze(1)
        hm = hm.to(context.dtype)
        if attn_probs is not None:
            attn_probs = attn_probs * hm
        context = context * hm

    # 6. head merge & output projection
    context = context.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
    context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

    if self.is_nested_encoder:
        output = context
    else:
        output = self.out_proj(context)

    # SVDLA-FIX-06: explicit error on shape mismatch
    if output.shape[-1] != hidden_states.shape[-1]:
        raise RuntimeError(
            f"SVDLinearAttention: output dim {output.shape[-1]} != "
            f"input hidden dim {hidden_states.shape[-1]}"
        )

    attn_out = attn_probs if output_attentions else None

    if return_state and updated_state is not None:
        return output, attn_out, updated_state
    return output, attn_out

Must return (output, attention_weights_or_None).

def get_orthogonality_loss(self) ‑> torch.Tensor

Expand source code

def get_orthogonality_loss(self) -> torch.Tensor:
    return (
        self.q_proj.get_orthogonality_loss()
        + self.k_proj.get_orthogonality_loss()
        + self.v_proj.get_orthogonality_loss()
        + self.out_proj.get_orthogonality_loss()
    )

def step_recurrent(self,
x_t: torch.Tensor,
past_state: tuple[torch.Tensor, torch.Tensor] | None = None) ‑> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]

Expand source code

def step_recurrent(
    self,
    x_t: torch.Tensor,
    past_state: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    if x_t.dim() == 2:
        x_t = x_t.unsqueeze(1)

    B = x_t.shape[0]
    q = self.q_proj(x_t).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
    k = (
        self.k_proj(x_t)
        .view(B, 1, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v = (
        self.v_proj(x_t)
        .view(B, 1, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    if self.num_kv_heads < self.num_heads:
        k = repeat_kv(k, self.kv_repeat)
        v = repeat_kv(v, self.kv_repeat)

    phi_q = F.elu(q.float()) + 1.0
    phi_k = F.elu(k.float()) + 1.0

    _, H, _, D_k = q.shape
    D_v = v.shape[-1]

    if self.use_decay_gating:
        decay_logits = self.decay_gate(x_t)
        if self.decay_channel_wise:
            decay_logits = decay_logits.view(B, 1, H, D_k).permute(0, 2, 1, 3)
            lam_t = (
                compute_exponential_decay(decay_logits, activation=self.decay_activation)
                .unsqueeze(-1)
                .squeeze(2)
            )
        else:
            decay_logits = decay_logits.permute(0, 2, 1).unsqueeze(-1).unsqueeze(-1)
            lam_t = compute_exponential_decay(
                decay_logits, activation=self.decay_activation
            ).squeeze(2)
    else:
        lam_t = torch.ones(B, H, 1, 1, device=x_t.device, dtype=torch.float32)

    if past_state is not None:
        S_prev, Z_prev = past_state
    else:
        S_prev = torch.zeros(B, H, D_k, D_v, device=x_t.device, dtype=torch.float32)
        Z_prev = torch.zeros(B, H, D_k, 1, device=x_t.device, dtype=torch.float32)

    kt_t = phi_k.transpose(-1, -2).contiguous()
    kv_t = torch.matmul(kt_t, v.float())

    S_t = lam_t * S_prev + kv_t
    Z_t = lam_t * Z_prev + kt_t

    num_t = torch.matmul(phi_q, S_t)
    den_t = torch.matmul(phi_q, Z_t) + self.eps
    y_t = num_t / den_t

    context = y_t.transpose(1, 2).contiguous().view(B, 1, self.all_head_size).to(x_t.dtype)
    if self.is_nested_encoder:
        out_t = context
    else:
        out_t = self.out_proj(context)

    if out_t.shape[-1] != x_t.shape[-1]:
        raise RuntimeError(
            f"SVDLinearAttention.step_recurrent: output dim "
            f"{out_t.shape[-1]} != input dim {x_t.shape[-1]}"
        )

    return out_t, (S_t, Z_t)

class LowRankSVDSequenceAttention (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = 64,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
dropout: float | None = 0.1,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("svd_seq", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_sequence", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_sequence_attention", supports_causal=True, supports_gqa=True)
class SVDLinearAttention(BaseAttention):
    r"""Low-Rank SVD-Compressed Sequence Attention Mechanism.
    Hardened for production use across diverse transformer architectures.
    """

    supports_gqa: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "low_rank"

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = 64,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        dropout: float | None = 0.1,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        # 1. Dimensions & Architecture
        self.hidden_size: int = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.num_heads: int = int(
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "num_heads", 12),
            )
        )
        self.head_dim: int = int(
            _cfg_get(
                config,
                "head_dim",
                kwargs.get(
                    "head_dim",
                    self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
                ),
            )
        )

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = max(1, self.num_heads // self.num_kv_heads)
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.q_head_size: int = self.num_heads * self.head_dim
        self.all_head_size: int = self.q_head_size

        # Explicitly allow None to disable SVD compression
        if rank is None:
            svd_disabled = _cfg_get(config, "svd_disabled", False)
            if svd_disabled or ("rank" in kwargs and kwargs["rank"] is None):
                self.rank: int | None = None
            else:
                cfg_r = _cfg_get(config, "svd_rank", _cfg_get(config, "rank", None))
                self.rank = int(cfg_r) if cfg_r is not None else None
        else:
            self.rank = int(rank)
        self.lora_rank: int = int(lora_rank if lora_rank is not None else 16)
        self.lora_alpha: float = lora_alpha

        # 2. Causal decoder flag detection
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = (
                bool(_cfg_get(config, "is_decoder", False))
                or any(
                    str(_cfg_get(config, "model_type", "")).lower().startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                or (old_attn is not None and hasattr(old_attn, "c_attn"))
            )

        self.is_nested_encoder: bool = False

        # 3. Projections with LoRA
        self.query = LoRALinear(
            self.hidden_size,
            self.q_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.key = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.value = 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.q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_val: float = dropout if dropout is not None else getattr(config, "dropout", 0.1)
        self.dropout: nn.Dropout = nn.Dropout(dropout_val)
        self._last_u_r: torch.Tensor | None = None

        # Persistent SVD projection buffer (populated in copy_weights from W_K)
        # Shape: [hidden_size, rank] — the top-r right singular vectors of W_K
        init_rank = max(1, self.rank or 64)
        self.register_buffer("svd_proj", torch.eye(self.hidden_size, init_rank), persistent=True)
        self._svd_proj_ready: bool = False
        self.triton_kwargs = {}

        if old_attn is not None:
            self.copy_weights(old_attn)

    @property
    def q_proj(self) -> LoRALinear:
        return self.query

    @property
    def k_proj(self) -> LoRALinear:
        return self.key

    @property
    def v_proj(self) -> LoRALinear:
        return self.value

    @property
    def q_lin(self) -> LoRALinear:
        return self.query

    @property
    def k_lin(self) -> LoRALinear:
        return self.key

    @property
    def v_lin(self) -> LoRALinear:
        return self.value

    def copy_weights(self, old_attn: nn.Module) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Detect nested encoder container or inner self-attention
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                # Unwrap nested encoder modules to find innermost attention
                src = old_attn
                while hasattr(src, "self") and not any(
                    hasattr(src, attr)
                    for attr in ["query", "key", "value", "q_proj", "k_proj", "v_proj"]
                ):
                    src = src.self
                # Copy output projection if available
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                p_cls = getattr(src.c_attn, "__class__", type(src.c_attn)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(src.c_attn, "nf")
                    or (
                        c_attn_w.shape[0] == self.hidden_size
                        and c_attn_w.shape[1] != self.hidden_size
                    )
                ):
                    c_attn_w = c_attn_w.t().contiguous()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size

                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

                out_w = src.c_proj.weight.detach()
                p_proj_cls = getattr(src.c_proj, "__class__", type(src.c_proj)).__name__
                if (
                    ("Conv1D" in p_proj_cls)
                    or hasattr(src.c_proj, "nf")
                    or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
                ):
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # Fallback: UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    w_q, b_q = UniversalAdapter.get_weight_bias(old_attn, "q")
                    w_k, b_k = UniversalAdapter.get_weight_bias(old_attn, "k")
                    w_v, b_v = UniversalAdapter.get_weight_bias(old_attn, "v")
                    if w_o is None:
                        w_o, b_o = UniversalAdapter.get_weight_bias(old_attn, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # Verification of essential projections
            missing = []
            if w_q is None:
                missing.append("query/q_proj")
            if w_k is None:
                missing.append("key/k_proj")
            if w_v is None:
                missing.append("value/v_proj")
            if missing:
                raise RuntimeError(
                    f"SVD-Compressed Sequence Attention: cannot extract "
                    f"{missing} from {old_attn.__class__.__name__}. "
                    f"Aborting to avoid silent corruption."
                )

            # Frobenius norm checks with None guard
            for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
                if w is None:
                    continue
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    raise RuntimeError(
                        f"SVD-Sequence: {name} weight has near-zero "
                        f"Frobenius norm. This will destroy the model."
                    )

            # Copy Q
            if w_q.shape == self.query.weight.shape:
                self.query.weight.copy_(w_q)
            else:
                self.query.weight.data[
                    : min(self.q_head_size, w_q.shape[0]),
                    : min(self.hidden_size, w_q.shape[1]),
                ].copy_(
                    w_q[
                        : min(self.q_head_size, w_q.shape[0]),
                        : min(self.hidden_size, w_q.shape[1]),
                    ]
                )
            if b_q is not None and self.query.bias is not None:
                self.query.bias.data[: min(len(self.query.bias), len(b_q))].copy_(
                    b_q[: min(len(self.query.bias), len(b_q))]
                )

            # Copy K & V with GQA pooling/repetition if needed
            def _copy_kv(target_proj, w, b):
                if w.shape[0] == self.kv_head_size:
                    target_proj.weight.copy_(w)
                    if b is not None and target_proj.bias is not None:
                        target_proj.bias.copy_(b)
                elif w.shape[0] > self.kv_head_size and w.shape[0] % self.kv_head_size == 0:
                    # GQA: take the first group from each cluster (deterministic)
                    rep = w.shape[0] // self.kv_head_size
                    w_grouped = w.view(self.num_kv_heads, rep, self.head_dim, -1)[:, 0].reshape(
                        self.kv_head_size, -1
                    )
                    target_proj.weight.copy_(w_grouped)
                    if b is not None and target_proj.bias is not None:
                        b_grouped = b.view(self.num_kv_heads, rep, self.head_dim)[:, 0].reshape(-1)
                        target_proj.bias.copy_(b_grouped)
                elif w.shape[0] < self.kv_head_size and self.kv_head_size % w.shape[0] == 0:
                    rep = self.kv_head_size // w.shape[0]
                    target_proj.weight.copy_(w.repeat(rep, 1))
                    if b is not None and target_proj.bias is not None:
                        target_proj.bias.copy_(b.repeat(rep))
                else:
                    target_proj.weight.data[
                        : min(self.kv_head_size, w.shape[0]),
                        : min(self.hidden_size, w.shape[1]),
                    ].copy_(
                        w[
                            : min(self.kv_head_size, w.shape[0]),
                            : min(self.hidden_size, w.shape[1]),
                        ]
                    )
                    if b is not None and target_proj.bias is not None:
                        target_proj.bias.data[: min(len(target_proj.bias), len(b))].copy_(
                            b[: min(len(target_proj.bias), len(b))]
                        )

            _copy_kv(self.key, w_k, b_k)
            _copy_kv(self.value, w_v, b_v)

            # Always copy Out projection — our module fully replaces BertAttention
            # (including BertSelfOutput), so we must apply out_proj in all cases.
            if w_o is not None:
                if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                    logger.warning("out_proj weight has near-zero Frobenius norm — skipping copy")
                else:
                    if w_o.shape == self.out_proj.weight.shape:
                        self.out_proj.weight.copy_(w_o)
                    else:
                        r = min(self.out_proj.weight.shape[0], w_o.shape[0])
                        c = min(self.out_proj.weight.shape[1], w_o.shape[1])
                        self.out_proj.weight.data[:r, :c].copy_(w_o[:r, :c])
                    if b_o is not None and self.out_proj.bias is not None:
                        mn = min(len(self.out_proj.bias), len(b_o))
                        self.out_proj.bias.data[:mn].copy_(b_o[:mn])

            # Compute SVD of W_K to get a fixed weight-space projection
            # W_K shape: [kv_head_size, hidden_size] → SVD gives right singular vectors of shape [hidden_size, ...]
            # We use the top-r right singular vectors (Vh rows) as the projection basis
            try:
                with torch.no_grad():
                    wk = self.key.weight.float()  # [kv_head_size, hidden_size]
                    _, _, Vh = torch.linalg.svd(wk, full_matrices=False)
                    # Vh: [min(kv_head_size, hidden_size), hidden_size]
                    r = min(self.rank or 64, Vh.shape[0], self.hidden_size)
                    # svd_proj: [hidden_size, r] — project hidden_states → rank-r space
                    proj = Vh[:r, :].T  # [hidden_size, r]
                    if proj.shape == self.svd_proj.shape:
                        self.svd_proj.copy_(proj)
                    else:
                        mn_r = min(proj.shape[1], self.svd_proj.shape[1])
                        self.svd_proj[:, :mn_r].copy_(proj[:, :mn_r])
                    self._svd_proj_ready = True
                    logger.info(
                        "SVD projection computed from W_K: shape=%s, rank=%d",
                        self.svd_proj.shape,
                        r,
                    )
            except Exception as e:
                logger.warning(
                    "SVD projection computation failed: %s — falling back to identity",
                    e,
                )
                self._svd_proj_ready = False

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, seq_len, _ = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype
        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Project Q, K, V
        Q = (
            self.query(hidden_states)
            .view(batch_size, seq_len, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            self.key(kv_input)
            .view(batch_size, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V = (
            self.value(kv_input)
            .view(batch_size, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. RoPE support
        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
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # 3. GQA expansion
        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        is_causal = self.is_decoder and not is_cross

        # Full attention if sequence is too short for SVD compression
        # (SVD with r >= N is meaningless — the projection is identity)
        use_full_attention = (
            self.rank is None or kv_len <= (self.rank or 0) + 8  # small buffer
        )

        if use_full_attention:
            # Use SDPA when possible (no explicit head_mask/output_attentions)
            _use_sdpa = head_mask is None and output_attentions is False

            if _use_sdpa:
                try:
                    _q = Q if Q.dtype in (torch.float16, torch.bfloat16) else Q.half()
                    _k = K if K.dtype in (torch.float16, torch.bfloat16) else K.half()
                    _v = V if V.dtype in (torch.float16, torch.bfloat16) else V.half()

                    # Prepare attn_mask for SDPA if eff_mask is provided
                    _attn_mask = None
                    if eff_mask is not None and eff_mask.dim() == 2:
                        _attn_mask = eff_mask.bool().view(
                            eff_mask.shape[0], 1, 1, eff_mask.shape[1]
                        )
                    elif eff_mask is not None and eff_mask.dim() == 4:
                        _attn_mask = eff_mask.bool() if eff_mask.dtype == torch.bool else eff_mask

                    context = torch.nn.functional.scaled_dot_product_attention(
                        _q,
                        _k,
                        _v,
                        attn_mask=_attn_mask,
                        is_causal=is_causal if _attn_mask is None else False,
                        dropout_p=self.dropout.p if self.training else 0.0,
                    ).to(dtype)
                    probs = None  # not available in SDPA
                except Exception as e:
                    logger.debug("SDPA failed: %s, falling back", e)
                    _use_sdpa = False

            if not _use_sdpa:
                # Manual softmax path
                scores = torch.matmul(
                    Q.float(), K.float().transpose(-2, -1).contiguous()
                ) / math.sqrt(self.head_dim)
                if is_causal and seq_len > 1:
                    causal_mask = (
                        torch.triu(
                            torch.full(
                                (seq_len, kv_len),
                                float("-inf"),
                                device=device,
                                dtype=torch.float32,
                            ),
                            diagonal=1,
                        )
                        .unsqueeze(0)
                        .unsqueeze(0)
                    )
                    scores = scores + causal_mask
                if eff_mask is not None:
                    scores = apply_attention_mask(scores, eff_mask)
                scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
                probs = F.softmax(scores, dim=-1).to(dtype)
                probs = self.dropout(probs)
                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)
                    elif head_mask.dim() == 3:
                        head_mask = head_mask.unsqueeze(-1)
                    probs = probs * head_mask
                context = torch.matmul(probs, V)

            if attention_mask is not None:
                q_valid_mask = get_token_valid_mask(
                    attention_mask,
                    batch_size,
                    seq_len,
                    device,
                )
                if q_valid_mask is not None:
                    context = context * q_valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
            attn_ret = probs if output_attentions else None
        else:
            # Use the fixed SVD subspace precomputed from W_K in copy_weights.
            # This is stable across batches — LoRA can converge on a fixed subspace.
            # svd_proj: [hidden_size, r]  (registered buffer, frozen)
            with torch.no_grad():
                proj = self.svd_proj.to(device=device, dtype=dtype)  # [hidden_size, r]
                actual_r = proj.shape[1]

                # Project the kv_input sequence through the fixed subspace
                # kv_input: [B, kv_len, hidden_size] → scores: [B, kv_len, actual_r]
                seq_scores = torch.matmul(kv_input.float(), proj.float())  # [B, kv_len, r]
                try:
                    U_r, _ = torch.linalg.qr(seq_scores.float(), mode="reduced")  # [B, kv_len, r]
                except Exception:
                    U_r, _ = torch.linalg.qr(seq_scores.cpu().float(), mode="reduced")
                    U_r = U_r.to(device=device)
                U_r = U_r.to(dtype)  # [B, kv_len, actual_r]

            # Store detached — the ortho loss is a metric, not a gradient source
            self._last_u_r = U_r.detach()

            # Project K and V through the stable subspace
            U_r_t = U_r.unsqueeze(1)  # [B, 1, kv_len, actual_r]
            K_proj = torch.matmul(
                U_r_t.transpose(-2, -1).contiguous(), K
            )  # [B, H, actual_r, head_dim]
            V_proj = torch.matmul(
                U_r_t.transpose(-2, -1).contiguous(), V
            )  # [B, H, actual_r, head_dim]

            scores = torch.matmul(
                Q.float(), K_proj.float().transpose(-2, -1).contiguous()
            ) / math.sqrt(self.head_dim)

            if is_causal and seq_len > 1 and actual_r == kv_len:
                # Only apply causal mask if we kept the full sequence
                causal_mask = (
                    torch.triu(
                        torch.full(
                            (seq_len, actual_r),
                            float("-inf"),
                            device=device,
                            dtype=torch.float32,
                        ),
                        diagonal=1,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal_mask
            elif is_causal:
                logger.debug(
                    "SVD-Sequence: causal mode with r=%d < kv_len=%d. "
                    "Causal masking not strictly enforced because SVD "
                    "compression combines multiple tokens.",
                    actual_r,
                    kv_len,
                )

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(dtype)
            probs = self.dropout(probs)

            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)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(-1)
                probs = probs * head_mask

            context = torch.matmul(probs, V_proj)  # [B, H, seq_len, head_dim]
            if attention_mask is not None:
                q_valid_mask = get_token_valid_mask(
                    attention_mask,
                    batch_size,
                    seq_len,
                    device,
                )
                if q_valid_mask is not None:
                    context = context * q_valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
            attn_ret = probs if output_attentions else None

        # 4. Context merge & Output projection
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # Always apply out_proj — our module replaces the full BertAttention container
        out = self.out_proj(context)

        if out.shape[-1] != hidden_states.shape[-1]:
            if out.shape[-1] < hidden_states.shape[-1]:
                out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
            else:
                out = out[..., : hidden_states.shape[-1]]

        return (out, attn_ret)

    def get_ortho_loss(self) -> torch.Tensor:
        """
        Orthogonality diagnostic for the QR basis.
        Returns a scalar (no gradient). Use only for logging.
        """
        device = self.query.weight.device
        if self._last_u_r is None:
            return torch.tensor(0.0, device=device)
        U = self._last_u_r
        r = U.size(-1)
        # Mean per-batch orthogonality error
        gram = torch.matmul(U.transpose(-2, -1).float(), U.float())
        I = torch.eye(r, device=U.device, dtype=torch.float32)
        # Average over batch
        loss = (gram - I.unsqueeze(0)).abs().mean()
        return loss.detach()

Low-Rank SVD-Compressed Sequence Attention Mechanism. Hardened for production use across diverse transformer architectures.

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_gqa : bool

The type of the None singleton.

var triton_kernel : str

The type of the None singleton.

Instance variables

prop k_lin : LoRALinear

Expand source code

@property
def k_lin(self) -> LoRALinear:
    return self.key

prop k_proj : LoRALinear

Expand source code

@property
def k_proj(self) -> LoRALinear:
    return self.key

prop q_lin : LoRALinear

Expand source code

@property
def q_lin(self) -> LoRALinear:
    return self.query

prop q_proj : LoRALinear

Expand source code

@property
def q_proj(self) -> LoRALinear:
    return self.query

prop v_lin : LoRALinear

Expand source code

@property
def v_lin(self) -> LoRALinear:
    return self.value

prop v_proj : LoRALinear

Expand source code

@property
def v_proj(self) -> LoRALinear:
    return self.value

Methods

def copy_weights(self, old_attn: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Detect nested encoder container or inner self-attention
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            # Unwrap nested encoder modules to find innermost attention
            src = old_attn
            while hasattr(src, "self") and not any(
                hasattr(src, attr)
                for attr in ["query", "key", "value", "q_proj", "k_proj", "v_proj"]
            ):
                src = src.self
            # Copy output projection if available
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight
                b_o = old_attn.output.dense.bias
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
            p_cls = getattr(src.c_attn, "__class__", type(src.c_attn)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(src.c_attn, "nf")
                or (
                    c_attn_w.shape[0] == self.hidden_size
                    and c_attn_w.shape[1] != self.hidden_size
                )
            ):
                c_attn_w = c_attn_w.t().contiguous()

            q_dim = self.q_head_size
            k_dim = self.kv_head_size
            v_dim = self.kv_head_size

            w_q = c_attn_w[:q_dim, :]
            w_k = c_attn_w[q_dim : q_dim + k_dim, :]
            w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

            if c_attn_b is not None:
                b_q = c_attn_b[:q_dim]
                b_k = c_attn_b[q_dim : q_dim + k_dim]
                b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

            out_w = src.c_proj.weight.detach()
            p_proj_cls = getattr(src.c_proj, "__class__", type(src.c_proj)).__name__
            if (
                ("Conv1D" in p_proj_cls)
                or hasattr(src.c_proj, "nf")
                or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
            ):
                out_w = out_w.t().contiguous()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # Fallback: UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                w_q, b_q = UniversalAdapter.get_weight_bias(old_attn, "q")
                w_k, b_k = UniversalAdapter.get_weight_bias(old_attn, "k")
                w_v, b_v = UniversalAdapter.get_weight_bias(old_attn, "v")
                if w_o is None:
                    w_o, b_o = UniversalAdapter.get_weight_bias(old_attn, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # Verification of essential projections
        missing = []
        if w_q is None:
            missing.append("query/q_proj")
        if w_k is None:
            missing.append("key/k_proj")
        if w_v is None:
            missing.append("value/v_proj")
        if missing:
            raise RuntimeError(
                f"SVD-Compressed Sequence Attention: cannot extract "
                f"{missing} from {old_attn.__class__.__name__}. "
                f"Aborting to avoid silent corruption."
            )

        # Frobenius norm checks with None guard
        for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
            if w is None:
                continue
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                raise RuntimeError(
                    f"SVD-Sequence: {name} weight has near-zero "
                    f"Frobenius norm. This will destroy the model."
                )

        # Copy Q
        if w_q.shape == self.query.weight.shape:
            self.query.weight.copy_(w_q)
        else:
            self.query.weight.data[
                : min(self.q_head_size, w_q.shape[0]),
                : min(self.hidden_size, w_q.shape[1]),
            ].copy_(
                w_q[
                    : min(self.q_head_size, w_q.shape[0]),
                    : min(self.hidden_size, w_q.shape[1]),
                ]
            )
        if b_q is not None and self.query.bias is not None:
            self.query.bias.data[: min(len(self.query.bias), len(b_q))].copy_(
                b_q[: min(len(self.query.bias), len(b_q))]
            )

        # Copy K & V with GQA pooling/repetition if needed
        def _copy_kv(target_proj, w, b):
            if w.shape[0] == self.kv_head_size:
                target_proj.weight.copy_(w)
                if b is not None and target_proj.bias is not None:
                    target_proj.bias.copy_(b)
            elif w.shape[0] > self.kv_head_size and w.shape[0] % self.kv_head_size == 0:
                # GQA: take the first group from each cluster (deterministic)
                rep = w.shape[0] // self.kv_head_size
                w_grouped = w.view(self.num_kv_heads, rep, self.head_dim, -1)[:, 0].reshape(
                    self.kv_head_size, -1
                )
                target_proj.weight.copy_(w_grouped)
                if b is not None and target_proj.bias is not None:
                    b_grouped = b.view(self.num_kv_heads, rep, self.head_dim)[:, 0].reshape(-1)
                    target_proj.bias.copy_(b_grouped)
            elif w.shape[0] < self.kv_head_size and self.kv_head_size % w.shape[0] == 0:
                rep = self.kv_head_size // w.shape[0]
                target_proj.weight.copy_(w.repeat(rep, 1))
                if b is not None and target_proj.bias is not None:
                    target_proj.bias.copy_(b.repeat(rep))
            else:
                target_proj.weight.data[
                    : min(self.kv_head_size, w.shape[0]),
                    : min(self.hidden_size, w.shape[1]),
                ].copy_(
                    w[
                        : min(self.kv_head_size, w.shape[0]),
                        : min(self.hidden_size, w.shape[1]),
                    ]
                )
                if b is not None and target_proj.bias is not None:
                    target_proj.bias.data[: min(len(target_proj.bias), len(b))].copy_(
                        b[: min(len(target_proj.bias), len(b))]
                    )

        _copy_kv(self.key, w_k, b_k)
        _copy_kv(self.value, w_v, b_v)

        # Always copy Out projection — our module fully replaces BertAttention
        # (including BertSelfOutput), so we must apply out_proj in all cases.
        if w_o is not None:
            if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                logger.warning("out_proj weight has near-zero Frobenius norm — skipping copy")
            else:
                if w_o.shape == self.out_proj.weight.shape:
                    self.out_proj.weight.copy_(w_o)
                else:
                    r = min(self.out_proj.weight.shape[0], w_o.shape[0])
                    c = min(self.out_proj.weight.shape[1], w_o.shape[1])
                    self.out_proj.weight.data[:r, :c].copy_(w_o[:r, :c])
                if b_o is not None and self.out_proj.bias is not None:
                    mn = min(len(self.out_proj.bias), len(b_o))
                    self.out_proj.bias.data[:mn].copy_(b_o[:mn])

        # Compute SVD of W_K to get a fixed weight-space projection
        # W_K shape: [kv_head_size, hidden_size] → SVD gives right singular vectors of shape [hidden_size, ...]
        # We use the top-r right singular vectors (Vh rows) as the projection basis
        try:
            with torch.no_grad():
                wk = self.key.weight.float()  # [kv_head_size, hidden_size]
                _, _, Vh = torch.linalg.svd(wk, full_matrices=False)
                # Vh: [min(kv_head_size, hidden_size), hidden_size]
                r = min(self.rank or 64, Vh.shape[0], self.hidden_size)
                # svd_proj: [hidden_size, r] — project hidden_states → rank-r space
                proj = Vh[:r, :].T  # [hidden_size, r]
                if proj.shape == self.svd_proj.shape:
                    self.svd_proj.copy_(proj)
                else:
                    mn_r = min(proj.shape[1], self.svd_proj.shape[1])
                    self.svd_proj[:, :mn_r].copy_(proj[:, :mn_r])
                self._svd_proj_ready = True
                logger.info(
                    "SVD projection computed from W_K: shape=%s, rank=%d",
                    self.svd_proj.shape,
                    r,
                )
        except Exception as e:
            logger.warning(
                "SVD projection computation failed: %s — falling back to identity",
                e,
            )
            self._svd_proj_ready = False

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

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,
*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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    batch_size, seq_len, _ = hidden_states.shape
    device = hidden_states.device
    dtype = hidden_states.dtype
    is_cross = encoder_hidden_states is not None
    kv_input = encoder_hidden_states if is_cross else hidden_states
    kv_len = kv_input.shape[1]

    # 1. Project Q, K, V
    Q = (
        self.query(hidden_states)
        .view(batch_size, seq_len, self.num_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    K = (
        self.key(kv_input)
        .view(batch_size, kv_len, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    V = (
        self.value(kv_input)
        .view(batch_size, kv_len, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # 2. RoPE support
    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
        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

    # 3. GQA expansion
    if self.num_kv_heads < self.num_heads:
        K = repeat_kv(K, self.kv_repeat)
        V = repeat_kv(V, self.kv_repeat)

    eff_mask = encoder_attention_mask if is_cross else attention_mask
    is_causal = self.is_decoder and not is_cross

    # Full attention if sequence is too short for SVD compression
    # (SVD with r >= N is meaningless — the projection is identity)
    use_full_attention = (
        self.rank is None or kv_len <= (self.rank or 0) + 8  # small buffer
    )

    if use_full_attention:
        # Use SDPA when possible (no explicit head_mask/output_attentions)
        _use_sdpa = head_mask is None and output_attentions is False

        if _use_sdpa:
            try:
                _q = Q if Q.dtype in (torch.float16, torch.bfloat16) else Q.half()
                _k = K if K.dtype in (torch.float16, torch.bfloat16) else K.half()
                _v = V if V.dtype in (torch.float16, torch.bfloat16) else V.half()

                # Prepare attn_mask for SDPA if eff_mask is provided
                _attn_mask = None
                if eff_mask is not None and eff_mask.dim() == 2:
                    _attn_mask = eff_mask.bool().view(
                        eff_mask.shape[0], 1, 1, eff_mask.shape[1]
                    )
                elif eff_mask is not None and eff_mask.dim() == 4:
                    _attn_mask = eff_mask.bool() if eff_mask.dtype == torch.bool else eff_mask

                context = torch.nn.functional.scaled_dot_product_attention(
                    _q,
                    _k,
                    _v,
                    attn_mask=_attn_mask,
                    is_causal=is_causal if _attn_mask is None else False,
                    dropout_p=self.dropout.p if self.training else 0.0,
                ).to(dtype)
                probs = None  # not available in SDPA
            except Exception as e:
                logger.debug("SDPA failed: %s, falling back", e)
                _use_sdpa = False

        if not _use_sdpa:
            # Manual softmax path
            scores = torch.matmul(
                Q.float(), K.float().transpose(-2, -1).contiguous()
            ) / math.sqrt(self.head_dim)
            if is_causal and seq_len > 1:
                causal_mask = (
                    torch.triu(
                        torch.full(
                            (seq_len, kv_len),
                            float("-inf"),
                            device=device,
                            dtype=torch.float32,
                        ),
                        diagonal=1,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal_mask
            if eff_mask is not None:
                scores = apply_attention_mask(scores, eff_mask)
            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(dtype)
            probs = self.dropout(probs)
            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)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(-1)
                probs = probs * head_mask
            context = torch.matmul(probs, V)

        if attention_mask is not None:
            q_valid_mask = get_token_valid_mask(
                attention_mask,
                batch_size,
                seq_len,
                device,
            )
            if q_valid_mask is not None:
                context = context * q_valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
        attn_ret = probs if output_attentions else None
    else:
        # Use the fixed SVD subspace precomputed from W_K in copy_weights.
        # This is stable across batches — LoRA can converge on a fixed subspace.
        # svd_proj: [hidden_size, r]  (registered buffer, frozen)
        with torch.no_grad():
            proj = self.svd_proj.to(device=device, dtype=dtype)  # [hidden_size, r]
            actual_r = proj.shape[1]

            # Project the kv_input sequence through the fixed subspace
            # kv_input: [B, kv_len, hidden_size] → scores: [B, kv_len, actual_r]
            seq_scores = torch.matmul(kv_input.float(), proj.float())  # [B, kv_len, r]
            try:
                U_r, _ = torch.linalg.qr(seq_scores.float(), mode="reduced")  # [B, kv_len, r]
            except Exception:
                U_r, _ = torch.linalg.qr(seq_scores.cpu().float(), mode="reduced")
                U_r = U_r.to(device=device)
            U_r = U_r.to(dtype)  # [B, kv_len, actual_r]

        # Store detached — the ortho loss is a metric, not a gradient source
        self._last_u_r = U_r.detach()

        # Project K and V through the stable subspace
        U_r_t = U_r.unsqueeze(1)  # [B, 1, kv_len, actual_r]
        K_proj = torch.matmul(
            U_r_t.transpose(-2, -1).contiguous(), K
        )  # [B, H, actual_r, head_dim]
        V_proj = torch.matmul(
            U_r_t.transpose(-2, -1).contiguous(), V
        )  # [B, H, actual_r, head_dim]

        scores = torch.matmul(
            Q.float(), K_proj.float().transpose(-2, -1).contiguous()
        ) / math.sqrt(self.head_dim)

        if is_causal and seq_len > 1 and actual_r == kv_len:
            # Only apply causal mask if we kept the full sequence
            causal_mask = (
                torch.triu(
                    torch.full(
                        (seq_len, actual_r),
                        float("-inf"),
                        device=device,
                        dtype=torch.float32,
                    ),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask
        elif is_causal:
            logger.debug(
                "SVD-Sequence: causal mode with r=%d < kv_len=%d. "
                "Causal masking not strictly enforced because SVD "
                "compression combines multiple tokens.",
                actual_r,
                kv_len,
            )

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(dtype)
        probs = self.dropout(probs)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            probs = probs * head_mask

        context = torch.matmul(probs, V_proj)  # [B, H, seq_len, head_dim]
        if attention_mask is not None:
            q_valid_mask = get_token_valid_mask(
                attention_mask,
                batch_size,
                seq_len,
                device,
            )
            if q_valid_mask is not None:
                context = context * q_valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
        attn_ret = probs if output_attentions else None

    # 4. Context merge & Output projection
    context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.all_head_size)
    context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

    # Always apply out_proj — our module replaces the full BertAttention container
    out = self.out_proj(context)

    if out.shape[-1] != hidden_states.shape[-1]:
        if out.shape[-1] < hidden_states.shape[-1]:
            out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
        else:
            out = out[..., : hidden_states.shape[-1]]

    return (out, attn_ret)

Must return (output, attention_weights_or_None).

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    """
    Orthogonality diagnostic for the QR basis.
    Returns a scalar (no gradient). Use only for logging.
    """
    device = self.query.weight.device
    if self._last_u_r is None:
        return torch.tensor(0.0, device=device)
    U = self._last_u_r
    r = U.size(-1)
    # Mean per-batch orthogonality error
    gram = torch.matmul(U.transpose(-2, -1).float(), U.float())
    I = torch.eye(r, device=U.device, dtype=torch.float32)
    # Average over batch
    loss = (gram - I.unsqueeze(0)).abs().mean()
    return loss.detach()

Orthogonality diagnostic for the QR basis. Returns a scalar (no gradient). Use only for logging.

class SVDAttention (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = 64,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
dropout: float | None = 0.1,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("svd_seq", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_sequence", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_sequence_attention", supports_causal=True, supports_gqa=True)
class SVDLinearAttention(BaseAttention):
    r"""Low-Rank SVD-Compressed Sequence Attention Mechanism.
    Hardened for production use across diverse transformer architectures.
    """

    supports_gqa: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "low_rank"

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = 64,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        dropout: float | None = 0.1,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        # 1. Dimensions & Architecture
        self.hidden_size: int = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.num_heads: int = int(
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "num_heads", 12),
            )
        )
        self.head_dim: int = int(
            _cfg_get(
                config,
                "head_dim",
                kwargs.get(
                    "head_dim",
                    self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
                ),
            )
        )

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = max(1, self.num_heads // self.num_kv_heads)
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.q_head_size: int = self.num_heads * self.head_dim
        self.all_head_size: int = self.q_head_size

        # Explicitly allow None to disable SVD compression
        if rank is None:
            svd_disabled = _cfg_get(config, "svd_disabled", False)
            if svd_disabled or ("rank" in kwargs and kwargs["rank"] is None):
                self.rank: int | None = None
            else:
                cfg_r = _cfg_get(config, "svd_rank", _cfg_get(config, "rank", None))
                self.rank = int(cfg_r) if cfg_r is not None else None
        else:
            self.rank = int(rank)
        self.lora_rank: int = int(lora_rank if lora_rank is not None else 16)
        self.lora_alpha: float = lora_alpha

        # 2. Causal decoder flag detection
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            self.is_decoder = (
                bool(_cfg_get(config, "is_decoder", False))
                or any(
                    str(_cfg_get(config, "model_type", "")).lower().startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
                or (old_attn is not None and hasattr(old_attn, "c_attn"))
            )

        self.is_nested_encoder: bool = False

        # 3. Projections with LoRA
        self.query = LoRALinear(
            self.hidden_size,
            self.q_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.key = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.value = 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.q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_val: float = dropout if dropout is not None else getattr(config, "dropout", 0.1)
        self.dropout: nn.Dropout = nn.Dropout(dropout_val)
        self._last_u_r: torch.Tensor | None = None

        # Persistent SVD projection buffer (populated in copy_weights from W_K)
        # Shape: [hidden_size, rank] — the top-r right singular vectors of W_K
        init_rank = max(1, self.rank or 64)
        self.register_buffer("svd_proj", torch.eye(self.hidden_size, init_rank), persistent=True)
        self._svd_proj_ready: bool = False
        self.triton_kwargs = {}

        if old_attn is not None:
            self.copy_weights(old_attn)

    @property
    def q_proj(self) -> LoRALinear:
        return self.query

    @property
    def k_proj(self) -> LoRALinear:
        return self.key

    @property
    def v_proj(self) -> LoRALinear:
        return self.value

    @property
    def q_lin(self) -> LoRALinear:
        return self.query

    @property
    def k_lin(self) -> LoRALinear:
        return self.key

    @property
    def v_lin(self) -> LoRALinear:
        return self.value

    def copy_weights(self, old_attn: nn.Module) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Detect nested encoder container or inner self-attention
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                # Unwrap nested encoder modules to find innermost attention
                src = old_attn
                while hasattr(src, "self") and not any(
                    hasattr(src, attr)
                    for attr in ["query", "key", "value", "q_proj", "k_proj", "v_proj"]
                ):
                    src = src.self
                # Copy output projection if available
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                p_cls = getattr(src.c_attn, "__class__", type(src.c_attn)).__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(src.c_attn, "nf")
                    or (
                        c_attn_w.shape[0] == self.hidden_size
                        and c_attn_w.shape[1] != self.hidden_size
                    )
                ):
                    c_attn_w = c_attn_w.t().contiguous()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size

                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

                out_w = src.c_proj.weight.detach()
                p_proj_cls = getattr(src.c_proj, "__class__", type(src.c_proj)).__name__
                if (
                    ("Conv1D" in p_proj_cls)
                    or hasattr(src.c_proj, "nf")
                    or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
                ):
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # Fallback: UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    w_q, b_q = UniversalAdapter.get_weight_bias(old_attn, "q")
                    w_k, b_k = UniversalAdapter.get_weight_bias(old_attn, "k")
                    w_v, b_v = UniversalAdapter.get_weight_bias(old_attn, "v")
                    if w_o is None:
                        w_o, b_o = UniversalAdapter.get_weight_bias(old_attn, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # Verification of essential projections
            missing = []
            if w_q is None:
                missing.append("query/q_proj")
            if w_k is None:
                missing.append("key/k_proj")
            if w_v is None:
                missing.append("value/v_proj")
            if missing:
                raise RuntimeError(
                    f"SVD-Compressed Sequence Attention: cannot extract "
                    f"{missing} from {old_attn.__class__.__name__}. "
                    f"Aborting to avoid silent corruption."
                )

            # Frobenius norm checks with None guard
            for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
                if w is None:
                    continue
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    raise RuntimeError(
                        f"SVD-Sequence: {name} weight has near-zero "
                        f"Frobenius norm. This will destroy the model."
                    )

            # Copy Q
            if w_q.shape == self.query.weight.shape:
                self.query.weight.copy_(w_q)
            else:
                self.query.weight.data[
                    : min(self.q_head_size, w_q.shape[0]),
                    : min(self.hidden_size, w_q.shape[1]),
                ].copy_(
                    w_q[
                        : min(self.q_head_size, w_q.shape[0]),
                        : min(self.hidden_size, w_q.shape[1]),
                    ]
                )
            if b_q is not None and self.query.bias is not None:
                self.query.bias.data[: min(len(self.query.bias), len(b_q))].copy_(
                    b_q[: min(len(self.query.bias), len(b_q))]
                )

            # Copy K & V with GQA pooling/repetition if needed
            def _copy_kv(target_proj, w, b):
                if w.shape[0] == self.kv_head_size:
                    target_proj.weight.copy_(w)
                    if b is not None and target_proj.bias is not None:
                        target_proj.bias.copy_(b)
                elif w.shape[0] > self.kv_head_size and w.shape[0] % self.kv_head_size == 0:
                    # GQA: take the first group from each cluster (deterministic)
                    rep = w.shape[0] // self.kv_head_size
                    w_grouped = w.view(self.num_kv_heads, rep, self.head_dim, -1)[:, 0].reshape(
                        self.kv_head_size, -1
                    )
                    target_proj.weight.copy_(w_grouped)
                    if b is not None and target_proj.bias is not None:
                        b_grouped = b.view(self.num_kv_heads, rep, self.head_dim)[:, 0].reshape(-1)
                        target_proj.bias.copy_(b_grouped)
                elif w.shape[0] < self.kv_head_size and self.kv_head_size % w.shape[0] == 0:
                    rep = self.kv_head_size // w.shape[0]
                    target_proj.weight.copy_(w.repeat(rep, 1))
                    if b is not None and target_proj.bias is not None:
                        target_proj.bias.copy_(b.repeat(rep))
                else:
                    target_proj.weight.data[
                        : min(self.kv_head_size, w.shape[0]),
                        : min(self.hidden_size, w.shape[1]),
                    ].copy_(
                        w[
                            : min(self.kv_head_size, w.shape[0]),
                            : min(self.hidden_size, w.shape[1]),
                        ]
                    )
                    if b is not None and target_proj.bias is not None:
                        target_proj.bias.data[: min(len(target_proj.bias), len(b))].copy_(
                            b[: min(len(target_proj.bias), len(b))]
                        )

            _copy_kv(self.key, w_k, b_k)
            _copy_kv(self.value, w_v, b_v)

            # Always copy Out projection — our module fully replaces BertAttention
            # (including BertSelfOutput), so we must apply out_proj in all cases.
            if w_o is not None:
                if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                    logger.warning("out_proj weight has near-zero Frobenius norm — skipping copy")
                else:
                    if w_o.shape == self.out_proj.weight.shape:
                        self.out_proj.weight.copy_(w_o)
                    else:
                        r = min(self.out_proj.weight.shape[0], w_o.shape[0])
                        c = min(self.out_proj.weight.shape[1], w_o.shape[1])
                        self.out_proj.weight.data[:r, :c].copy_(w_o[:r, :c])
                    if b_o is not None and self.out_proj.bias is not None:
                        mn = min(len(self.out_proj.bias), len(b_o))
                        self.out_proj.bias.data[:mn].copy_(b_o[:mn])

            # Compute SVD of W_K to get a fixed weight-space projection
            # W_K shape: [kv_head_size, hidden_size] → SVD gives right singular vectors of shape [hidden_size, ...]
            # We use the top-r right singular vectors (Vh rows) as the projection basis
            try:
                with torch.no_grad():
                    wk = self.key.weight.float()  # [kv_head_size, hidden_size]
                    _, _, Vh = torch.linalg.svd(wk, full_matrices=False)
                    # Vh: [min(kv_head_size, hidden_size), hidden_size]
                    r = min(self.rank or 64, Vh.shape[0], self.hidden_size)
                    # svd_proj: [hidden_size, r] — project hidden_states → rank-r space
                    proj = Vh[:r, :].T  # [hidden_size, r]
                    if proj.shape == self.svd_proj.shape:
                        self.svd_proj.copy_(proj)
                    else:
                        mn_r = min(proj.shape[1], self.svd_proj.shape[1])
                        self.svd_proj[:, :mn_r].copy_(proj[:, :mn_r])
                    self._svd_proj_ready = True
                    logger.info(
                        "SVD projection computed from W_K: shape=%s, rank=%d",
                        self.svd_proj.shape,
                        r,
                    )
            except Exception as e:
                logger.warning(
                    "SVD projection computation failed: %s — falling back to identity",
                    e,
                )
                self._svd_proj_ready = False

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, seq_len, _ = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype
        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Project Q, K, V
        Q = (
            self.query(hidden_states)
            .view(batch_size, seq_len, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            self.key(kv_input)
            .view(batch_size, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V = (
            self.value(kv_input)
            .view(batch_size, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. RoPE support
        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
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # 3. GQA expansion
        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        is_causal = self.is_decoder and not is_cross

        # Full attention if sequence is too short for SVD compression
        # (SVD with r >= N is meaningless — the projection is identity)
        use_full_attention = (
            self.rank is None or kv_len <= (self.rank or 0) + 8  # small buffer
        )

        if use_full_attention:
            # Use SDPA when possible (no explicit head_mask/output_attentions)
            _use_sdpa = head_mask is None and output_attentions is False

            if _use_sdpa:
                try:
                    _q = Q if Q.dtype in (torch.float16, torch.bfloat16) else Q.half()
                    _k = K if K.dtype in (torch.float16, torch.bfloat16) else K.half()
                    _v = V if V.dtype in (torch.float16, torch.bfloat16) else V.half()

                    # Prepare attn_mask for SDPA if eff_mask is provided
                    _attn_mask = None
                    if eff_mask is not None and eff_mask.dim() == 2:
                        _attn_mask = eff_mask.bool().view(
                            eff_mask.shape[0], 1, 1, eff_mask.shape[1]
                        )
                    elif eff_mask is not None and eff_mask.dim() == 4:
                        _attn_mask = eff_mask.bool() if eff_mask.dtype == torch.bool else eff_mask

                    context = torch.nn.functional.scaled_dot_product_attention(
                        _q,
                        _k,
                        _v,
                        attn_mask=_attn_mask,
                        is_causal=is_causal if _attn_mask is None else False,
                        dropout_p=self.dropout.p if self.training else 0.0,
                    ).to(dtype)
                    probs = None  # not available in SDPA
                except Exception as e:
                    logger.debug("SDPA failed: %s, falling back", e)
                    _use_sdpa = False

            if not _use_sdpa:
                # Manual softmax path
                scores = torch.matmul(
                    Q.float(), K.float().transpose(-2, -1).contiguous()
                ) / math.sqrt(self.head_dim)
                if is_causal and seq_len > 1:
                    causal_mask = (
                        torch.triu(
                            torch.full(
                                (seq_len, kv_len),
                                float("-inf"),
                                device=device,
                                dtype=torch.float32,
                            ),
                            diagonal=1,
                        )
                        .unsqueeze(0)
                        .unsqueeze(0)
                    )
                    scores = scores + causal_mask
                if eff_mask is not None:
                    scores = apply_attention_mask(scores, eff_mask)
                scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
                probs = F.softmax(scores, dim=-1).to(dtype)
                probs = self.dropout(probs)
                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)
                    elif head_mask.dim() == 3:
                        head_mask = head_mask.unsqueeze(-1)
                    probs = probs * head_mask
                context = torch.matmul(probs, V)

            if attention_mask is not None:
                q_valid_mask = get_token_valid_mask(
                    attention_mask,
                    batch_size,
                    seq_len,
                    device,
                )
                if q_valid_mask is not None:
                    context = context * q_valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
            attn_ret = probs if output_attentions else None
        else:
            # Use the fixed SVD subspace precomputed from W_K in copy_weights.
            # This is stable across batches — LoRA can converge on a fixed subspace.
            # svd_proj: [hidden_size, r]  (registered buffer, frozen)
            with torch.no_grad():
                proj = self.svd_proj.to(device=device, dtype=dtype)  # [hidden_size, r]
                actual_r = proj.shape[1]

                # Project the kv_input sequence through the fixed subspace
                # kv_input: [B, kv_len, hidden_size] → scores: [B, kv_len, actual_r]
                seq_scores = torch.matmul(kv_input.float(), proj.float())  # [B, kv_len, r]
                try:
                    U_r, _ = torch.linalg.qr(seq_scores.float(), mode="reduced")  # [B, kv_len, r]
                except Exception:
                    U_r, _ = torch.linalg.qr(seq_scores.cpu().float(), mode="reduced")
                    U_r = U_r.to(device=device)
                U_r = U_r.to(dtype)  # [B, kv_len, actual_r]

            # Store detached — the ortho loss is a metric, not a gradient source
            self._last_u_r = U_r.detach()

            # Project K and V through the stable subspace
            U_r_t = U_r.unsqueeze(1)  # [B, 1, kv_len, actual_r]
            K_proj = torch.matmul(
                U_r_t.transpose(-2, -1).contiguous(), K
            )  # [B, H, actual_r, head_dim]
            V_proj = torch.matmul(
                U_r_t.transpose(-2, -1).contiguous(), V
            )  # [B, H, actual_r, head_dim]

            scores = torch.matmul(
                Q.float(), K_proj.float().transpose(-2, -1).contiguous()
            ) / math.sqrt(self.head_dim)

            if is_causal and seq_len > 1 and actual_r == kv_len:
                # Only apply causal mask if we kept the full sequence
                causal_mask = (
                    torch.triu(
                        torch.full(
                            (seq_len, actual_r),
                            float("-inf"),
                            device=device,
                            dtype=torch.float32,
                        ),
                        diagonal=1,
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                scores = scores + causal_mask
            elif is_causal:
                logger.debug(
                    "SVD-Sequence: causal mode with r=%d < kv_len=%d. "
                    "Causal masking not strictly enforced because SVD "
                    "compression combines multiple tokens.",
                    actual_r,
                    kv_len,
                )

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(dtype)
            probs = self.dropout(probs)

            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)
                elif head_mask.dim() == 3:
                    head_mask = head_mask.unsqueeze(-1)
                probs = probs * head_mask

            context = torch.matmul(probs, V_proj)  # [B, H, seq_len, head_dim]
            if attention_mask is not None:
                q_valid_mask = get_token_valid_mask(
                    attention_mask,
                    batch_size,
                    seq_len,
                    device,
                )
                if q_valid_mask is not None:
                    context = context * q_valid_mask.unsqueeze(1).unsqueeze(-1).to(dtype)
            attn_ret = probs if output_attentions else None

        # 4. Context merge & Output projection
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.all_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # Always apply out_proj — our module replaces the full BertAttention container
        out = self.out_proj(context)

        if out.shape[-1] != hidden_states.shape[-1]:
            if out.shape[-1] < hidden_states.shape[-1]:
                out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
            else:
                out = out[..., : hidden_states.shape[-1]]

        return (out, attn_ret)

    def get_ortho_loss(self) -> torch.Tensor:
        """
        Orthogonality diagnostic for the QR basis.
        Returns a scalar (no gradient). Use only for logging.
        """
        device = self.query.weight.device
        if self._last_u_r is None:
            return torch.tensor(0.0, device=device)
        U = self._last_u_r
        r = U.size(-1)
        # Mean per-batch orthogonality error
        gram = torch.matmul(U.transpose(-2, -1).float(), U.float())
        I = torch.eye(r, device=U.device, dtype=torch.float32)
        # Average over batch
        loss = (gram - I.unsqueeze(0)).abs().mean()
        return loss.detach()

Low-Rank SVD-Compressed Sequence Attention Mechanism. Hardened for production use across diverse transformer architectures.

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_gqa : bool

The type of the None singleton.

var triton_kernel : str

The type of the None singleton.

Instance variables

prop k_lin : LoRALinear

Expand source code

@property
def k_lin(self) -> LoRALinear:
    return self.key

prop k_proj : LoRALinear

Expand source code

@property
def k_proj(self) -> LoRALinear:
    return self.key

prop q_lin : LoRALinear

Expand source code

@property
def q_lin(self) -> LoRALinear:
    return self.query

prop q_proj : LoRALinear

Expand source code

@property
def q_proj(self) -> LoRALinear:
    return self.query

prop v_lin : LoRALinear

Expand source code

@property
def v_lin(self) -> LoRALinear:
    return self.value

prop v_proj : LoRALinear

Expand source code

@property
def v_proj(self) -> LoRALinear:
    return self.value

Methods

def copy_weights(self, old_attn: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Detect nested encoder container or inner self-attention
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            # Unwrap nested encoder modules to find innermost attention
            src = old_attn
            while hasattr(src, "self") and not any(
                hasattr(src, attr)
                for attr in ["query", "key", "value", "q_proj", "k_proj", "v_proj"]
            ):
                src = src.self
            # Copy output projection if available
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight
                b_o = old_attn.output.dense.bias
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
            p_cls = getattr(src.c_attn, "__class__", type(src.c_attn)).__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(src.c_attn, "nf")
                or (
                    c_attn_w.shape[0] == self.hidden_size
                    and c_attn_w.shape[1] != self.hidden_size
                )
            ):
                c_attn_w = c_attn_w.t().contiguous()

            q_dim = self.q_head_size
            k_dim = self.kv_head_size
            v_dim = self.kv_head_size

            w_q = c_attn_w[:q_dim, :]
            w_k = c_attn_w[q_dim : q_dim + k_dim, :]
            w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

            if c_attn_b is not None:
                b_q = c_attn_b[:q_dim]
                b_k = c_attn_b[q_dim : q_dim + k_dim]
                b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

            out_w = src.c_proj.weight.detach()
            p_proj_cls = getattr(src.c_proj, "__class__", type(src.c_proj)).__name__
            if (
                ("Conv1D" in p_proj_cls)
                or hasattr(src.c_proj, "nf")
                or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
            ):
                out_w = out_w.t().contiguous()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # Fallback: UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                w_q, b_q = UniversalAdapter.get_weight_bias(old_attn, "q")
                w_k, b_k = UniversalAdapter.get_weight_bias(old_attn, "k")
                w_v, b_v = UniversalAdapter.get_weight_bias(old_attn, "v")
                if w_o is None:
                    w_o, b_o = UniversalAdapter.get_weight_bias(old_attn, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # Verification of essential projections
        missing = []
        if w_q is None:
            missing.append("query/q_proj")
        if w_k is None:
            missing.append("key/k_proj")
        if w_v is None:
            missing.append("value/v_proj")
        if missing:
            raise RuntimeError(
                f"SVD-Compressed Sequence Attention: cannot extract "
                f"{missing} from {old_attn.__class__.__name__}. "
                f"Aborting to avoid silent corruption."
            )

        # Frobenius norm checks with None guard
        for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
            if w is None:
                continue
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                raise RuntimeError(
                    f"SVD-Sequence: {name} weight has near-zero "
                    f"Frobenius norm. This will destroy the model."
                )

        # Copy Q
        if w_q.shape == self.query.weight.shape:
            self.query.weight.copy_(w_q)
        else:
            self.query.weight.data[
                : min(self.q_head_size, w_q.shape[0]),
                : min(self.hidden_size, w_q.shape[1]),
            ].copy_(
                w_q[
                    : min(self.q_head_size, w_q.shape[0]),
                    : min(self.hidden_size, w_q.shape[1]),
                ]
            )
        if b_q is not None and self.query.bias is not None:
            self.query.bias.data[: min(len(self.query.bias), len(b_q))].copy_(
                b_q[: min(len(self.query.bias), len(b_q))]
            )

        # Copy K & V with GQA pooling/repetition if needed
        def _copy_kv(target_proj, w, b):
            if w.shape[0] == self.kv_head_size:
                target_proj.weight.copy_(w)
                if b is not None and target_proj.bias is not None:
                    target_proj.bias.copy_(b)
            elif w.shape[0] > self.kv_head_size and w.shape[0] % self.kv_head_size == 0:
                # GQA: take the first group from each cluster (deterministic)
                rep = w.shape[0] // self.kv_head_size
                w_grouped = w.view(self.num_kv_heads, rep, self.head_dim, -1)[:, 0].reshape(
                    self.kv_head_size, -1
                )
                target_proj.weight.copy_(w_grouped)
                if b is not None and target_proj.bias is not None:
                    b_grouped = b.view(self.num_kv_heads, rep, self.head_dim)[:, 0].reshape(-1)
                    target_proj.bias.copy_(b_grouped)
            elif w.shape[0] < self.kv_head_size and self.kv_head_size % w.shape[0] == 0:
                rep = self.kv_head_size // w.shape[0]
                target_proj.weight.copy_(w.repeat(rep, 1))
                if b is not None and target_proj.bias is not None:
                    target_proj.bias.copy_(b.repeat(rep))
            else:
                target_proj.weight.data[
                    : min(self.kv_head_size, w.shape[0]),
                    : min(self.hidden_size, w.shape[1]),
                ].copy_(
                    w[
                        : min(self.kv_head_size, w.shape[0]),
                        : min(self.hidden_size, w.shape[1]),
                    ]
                )
                if b is not None and target_proj.bias is not None:
                    target_proj.bias.data[: min(len(target_proj.bias), len(b))].copy_(
                        b[: min(len(target_proj.bias), len(b))]
                    )

        _copy_kv(self.key, w_k, b_k)
        _copy_kv(self.value, w_v, b_v)

        # Always copy Out projection — our module fully replaces BertAttention
        # (including BertSelfOutput), so we must apply out_proj in all cases.
        if w_o is not None:
            if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                logger.warning("out_proj weight has near-zero Frobenius norm — skipping copy")
            else:
                if w_o.shape == self.out_proj.weight.shape:
                    self.out_proj.weight.copy_(w_o)
                else:
                    r = min(self.out_proj.weight.shape[0], w_o.shape[0])
                    c = min(self.out_proj.weight.shape[1], w_o.shape[1])
                    self.out_proj.weight.data[:r, :c].copy_(w_o[:r, :c])
                if b_o is not None and self.out_proj.bias is not None:
                    mn = min(len(self.out_proj.bias), len(b_o))
                    self.out_proj.bias.data[:mn].copy_(b_o[:mn])

        # Compute SVD of W_K to get a fixed weight-space projection
        # W_K shape: [kv_head_size, hidden_size] → SVD gives right singular vectors of shape [hidden_size, ...]
        # We use the top-r right singular vectors (Vh rows) as the projection basis
        try:
            with torch.no_grad():
                wk = self.key.weight.float()  # [kv_head_size, hidden_size]
                _, _, Vh = torch.linalg.svd(wk, full_matrices=False)
                # Vh: [min(kv_head_size, hidden_size), hidden_size]
                r = min(self.rank or 64, Vh.shape[0], self.hidden_size)
                # svd_proj: [hidden_size, r] — project hidden_states → rank-r space
                proj = Vh[:r, :].T  # [hidden_size, r]
                if proj.shape == self.svd_proj.shape:
                    self.svd_proj.copy_(proj)
                else:
                    mn_r = min(proj.shape[1], self.svd_proj.shape[1])
                    self.svd_proj[:, :mn_r].copy_(proj[:, :mn_r])
                self._svd_proj_ready = True
                logger.info(
                    "SVD projection computed from W_K: shape=%s, rank=%d",
                    self.svd_proj.shape,
                    r,
                )
        except Exception as e:
            logger.warning(
                "SVD projection computation failed: %s — falling back to identity",
                e,
            )
            self._svd_proj_ready = False

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

def get_ortho_loss(self) ‑> torch.Tensor

Expand source code

def get_ortho_loss(self) -> torch.Tensor:
    """
    Orthogonality diagnostic for the QR basis.
    Returns a scalar (no gradient). Use only for logging.
    """
    device = self.query.weight.device
    if self._last_u_r is None:
        return torch.tensor(0.0, device=device)
    U = self._last_u_r
    r = U.size(-1)
    # Mean per-batch orthogonality error
    gram = torch.matmul(U.transpose(-2, -1).float(), U.float())
    I = torch.eye(r, device=U.device, dtype=torch.float32)
    # Average over batch
    loss = (gram - I.unsqueeze(0)).abs().mean()
    return loss.detach()

Orthogonality diagnostic for the QR basis. Returns a scalar (no gradient). Use only for logging.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class SVDLinearProjection (in_features: int, out_features: int, rank: int = 64, bias: bool = True)

Expand source code

class SVDLinearProjection(nn.Module):
    r"""
    Genuine SVD Low-Rank Linear Layer: W \approx U @ V.

    Factorizes in_features -> out_features through bottleneck rank r:
      V \in R^{r x in_features}  (Down-projection)
      U \in R^{out_features x r} (Up-projection)
      y = (x @ V^T) @ U^T + bias
    """

    triton_kernel = "low_rank"

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 64,
        bias: bool = True,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = min(rank, in_features, out_features)

        self.V = nn.Parameter(torch.empty(self.rank, in_features))
        self.U = nn.Parameter(torch.empty(out_features, self.rank))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        # BUG 5 FIX: Identity-like init for U so initial output is non-zero before fine-tuning
        nn.init.kaiming_uniform_(self.V, a=math.sqrt(5))
        nn.init.zeros_(self.U)
        with torch.no_grad():
            r = min(self.out_features, self.rank)
            self.U[:r, :r].copy_(torch.eye(r))
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Fused low-rank projection: (x @ V^T) @ U^T + bias
        # Align weights with input dtype (e.g. under fp16/bf16 autocast) so Triton kernel runs directly
        u = self.U.to(x.dtype) if self.U.dtype != x.dtype else self.U
        v = self.V.to(x.dtype) if self.V.dtype != x.dtype else self.V
        b = (
            self.bias.to(x.dtype)
            if (self.bias is not None and self.bias.dtype != x.dtype)
            else self.bias
        )
        return triton_dispatch(self.triton_kernel, x, u, v, bias=b)

    @classmethod
    def from_pretrained_linear(cls, linear: nn.Linear, rank: int = 64) -> SVDLinearProjection:
        r"""
        Initializes low-rank factors directly from pretrained weights via Truncated SVD:
        W = U_full @ diag(S) @ V_full^T
        U = U_full[:, :r] @ sqrt(S[:r])
        V = sqrt(S[:r]) @ V_full^T[:r, :]
        """
        out_features, in_features = linear.weight.shape
        proj = cls(in_features, out_features, rank=rank, bias=(linear.bias is not None))

        with torch.no_grad():
            W = linear.weight.data.float()
            u, s, vh = torch.linalg.svd(W, full_matrices=False)

            r = proj.rank
            s_sqrt = torch.diag(torch.sqrt(s[:r]))
            proj.U.data.copy_(torch.matmul(u[:, :r], s_sqrt).to(linear.weight.dtype))
            proj.V.data.copy_(torch.matmul(s_sqrt, vh[:r, :]).to(linear.weight.dtype))

            if linear.bias is not None:
                proj.bias.data.copy_(linear.bias.data)

        return proj

Genuine SVD Low-Rank Linear Layer: W \approx U @ V.

Factorizes in_features -> out_features through bottleneck rank r: V \in R^{r x in_features} (Down-projection) U \in R^{out_features x r} (Up-projection) y = (x @ V^T) @ U^T + bias

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Class variables

var triton_kernel

The type of the None singleton.

Static methods

def from_pretrained_linear(linear: nn.Linear, rank: int = 64) ‑> SVDLinearProjection

Initializes low-rank factors directly from pretrained weights via Truncated SVD: W = U_full @ diag(S) @ V_full^T U = U_full[:, :r] @ sqrt(S[:r]) V = sqrt(S[:r]) @ V_full^T[:r, :]

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # Fused low-rank projection: (x @ V^T) @ U^T + bias
    # Align weights with input dtype (e.g. under fp16/bf16 autocast) so Triton kernel runs directly
    u = self.U.to(x.dtype) if self.U.dtype != x.dtype else self.U
    v = self.V.to(x.dtype) if self.V.dtype != x.dtype else self.V
    b = (
        self.bias.to(x.dtype)
        if (self.bias is not None and self.bias.dtype != x.dtype)
        else self.bias
    )
    return triton_dispatch(self.triton_kernel, x, u, v, bias=b)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    # BUG 5 FIX: Identity-like init for U so initial output is non-zero before fine-tuning
    nn.init.kaiming_uniform_(self.V, a=math.sqrt(5))
    nn.init.zeros_(self.U)
    with torch.no_grad():
        r = min(self.out_features, self.rank)
        self.U[:r, :r].copy_(torch.eye(r))
    if self.bias is not None:
        nn.init.zeros_(self.bias)

class SVDMamba (config: Any | None = None,
rank: int | None = None,
d_model: int | None = None,
d_state: int = 16,
d_conv: int = 4,
expand: int = 1,
dt_rank: int | None = None,
global_rank: int = 8,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("svd_mamba", supports_causal=True)
@ATTENTION_REGISTRY.register("SVD-Mamba")
@ATTENTION_REGISTRY.register("svd-mamba")
@ATTENTION_REGISTRY.register("SVD_Mamba")
class SVDMamba(BaseAttention):
    r"""SVD-Mamba module with SVD-factorized projections, causal convolution, and selective SSM."""

    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | None = None,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 1,
        dt_rank: int | None = None,
        global_rank: int = 8,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        comp_config = _cfg_get(config, "compression_config", kwargs.get("comp_config"))

        if d_model is None:
            if isinstance(config, int):
                d_model = config
            elif config is not None:
                d_model = _cfg_get(
                    config,
                    "hidden_size",
                    _cfg_get(config, "d_model", _cfg_get(config, "dim", 768)),
                )
        if d_model is None:
            d_model = kwargs.get("hidden_size", kwargs.get("d_model", 768))

        if hasattr(d_model, "hidden_size"):
            d_model = d_model.hidden_size

        self.d_model = int(d_model)
        self.d_state = _cfg_get(comp_config, "d_state", d_state)
        self.d_conv = _cfg_get(comp_config, "d_conv", d_conv)
        self.expand = _cfg_get(comp_config, "expand", expand)
        self.inner_dim = self.expand * self.d_model
        self.dt_rank = dt_rank if dt_rank is not None else max(1, math.ceil(self.d_model / 16))
        self.global_rank = min(global_rank, self.d_model, self.inner_dim)
        self.lora_rank = _cfg_get(
            config, "lora_rank", _cfg_get(comp_config, "lora_rank", lora_rank)
        )
        self.lora_alpha = _cfg_get(
            config, "lora_alpha", _cfg_get(comp_config, "lora_alpha", lora_alpha)
        )

        # Decoder / Causal flag
        model_type = str(_cfg_get(config, "model_type", "") or "").lower()
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_dec = False
        if is_decoder is not None:
            is_dec = is_decoder
        elif config is not None:
            is_dec = bool(_cfg_get(config, "is_decoder", False))
        else:
            is_dec = bool(kwargs.get("is_causal", False))
        if (
            any(model_type.startswith(p) for p in decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        self.is_decoder: bool = bool(is_dec)

        self.is_nested_encoder: bool = False
        num_heads = (
            int(_cfg_get(config, "num_attention_heads", kwargs.get("num_attention_heads", 12)))
            if config is not None
            else 12
        )
        num_kv_heads = (
            int(
                _cfg_get(
                    config,
                    "num_key_value_heads",
                    kwargs.get("num_key_value_heads", num_heads),
                )
            )
            if config is not None
            else num_heads
        )
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.head_dim = max(1, self.d_model // 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.supports_gqa = True

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        target_dtype = kwargs.get("dtype")
        if target_dtype is None and source_attn is not None:
            target_dtype = next(source_attn.parameters(), torch.tensor([])).dtype
            if not isinstance(target_dtype, torch.dtype):
                target_dtype = torch.float32

        # Attention path projections with LoRA
        self.q_proj = SVDLinear(
            self.d_model,
            self.all_head_size,
            rank=max(1, min(self.d_model, self.all_head_size) // 4),
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.k_proj = SVDLinear(
            self.d_model,
            self.kv_head_size,
            rank=max(1, min(self.d_model, self.kv_head_size) // 4),
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.v_proj = SVDLinear(
            self.d_model,
            self.kv_head_size,
            rank=max(1, min(self.d_model, self.kv_head_size) // 4),
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )

        # Gating for dynamic aggregation (negative initial bias preserves pretrained attention path)
        self.gating = nn.Linear(self.d_model, self.d_model)
        nn.init.zeros_(self.gating.weight)
        nn.init.constant_(self.gating.bias, -2.0)

        # Attention output projection
        self.attn_out_proj = SVDLinear(
            self.all_head_size,
            self.d_model,
            rank=max(1, min(self.all_head_size, self.d_model) // 2),
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )

        # SVD-factorized projections
        self.in_proj = SVDLinear(
            self.d_model,
            2 * self.inner_dim,
            rank=max(1, min(self.d_model, self.inner_dim) // 2),
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.out_proj = SVDLinear(
            self.inner_dim,
            self.d_model,
            rank=max(1, min(self.inner_dim, self.d_model) // 2),
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.x_proj = nn.Linear(self.inner_dim, self.dt_rank + 2 * self.d_state, bias=True)
        nn.init.uniform_(self.x_proj.weight, -0.01, 0.01)
        nn.init.zeros_(self.x_proj.bias)
        self.dt_proj = nn.Linear(self.dt_rank, self.inner_dim, bias=True)
        nn.init.uniform_(self.dt_proj.weight, -0.01, 0.01)
        nn.init.zeros_(self.dt_proj.bias)

        # Causal depthwise convolution
        self.conv1d = nn.Conv1d(
            in_channels=self.inner_dim,
            out_channels=self.inner_dim,
            kernel_size=self.d_conv,
            groups=self.inner_dim,
            bias=False,
        )

        # State-space parameters
        self.A_log = nn.Parameter(torch.log(torch.arange(1, self.d_state + 1, dtype=torch.float32)))
        self.D = nn.Parameter(torch.ones(self.inner_dim))
        self.dt_bias = nn.Parameter(torch.zeros(self.inner_dim))

        # Global SVD mixing
        self.global_V = nn.Parameter(torch.randn(self.d_model, self.global_rank) * 0.02)
        self.global_out = SVDLinear(
            self.global_rank,
            self.inner_dim,
            rank=self.global_rank,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=False,
            dtype=target_dtype,
        )

        if source_attn is not None:
            self.copy_weights(source_attn)

        if target_dtype is not None:
            self.to(target_dtype)

    @property
    def query(self) -> SVDLinear:
        return self.q_proj

    @property
    def key(self) -> SVDLinear:
        return self.k_proj

    @property
    def value(self) -> SVDLinear:
        return self.v_proj

    @property
    def q_lin(self) -> SVDLinear:
        return self.q_proj

    @property
    def k_lin(self) -> SVDLinear:
        return self.k_proj

    @property
    def v_lin(self) -> SVDLinear:
        return self.v_proj

    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Check nested encoder container
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o, b_o = _unpack_weight_bias(old_attn.output.dense)
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
                if c_attn_w is not None:
                    if (
                        c_attn_w.dim() == 2
                        and c_attn_w.shape[0] == self.d_model
                        and c_attn_w.shape[1] == 3 * self.d_model
                    ):
                        c_attn_w = c_attn_w.t().contiguous()
                    split_size = c_attn_w.size(0) // 3
                    w_q = c_attn_w[:split_size, :]
                    w_k = c_attn_w[split_size : 2 * split_size, :]
                    w_v = c_attn_w[2 * split_size :, :]
                    if c_attn_b is not None:
                        b_q = c_attn_b[:split_size]
                        b_k = c_attn_b[split_size : 2 * split_size]
                        b_v = c_attn_b[2 * split_size :]

                c_proj_w, c_proj_b = _unpack_weight_bias(src.c_proj)
                if c_proj_w is not None:
                    if (
                        c_proj_w.dim() == 2
                        and c_proj_w.shape[0] != self.d_model
                        and c_proj_w.shape[1] == self.d_model
                    ):
                        c_proj_w = c_proj_w.t().contiguous()
                    w_o = c_proj_w
                    b_o = c_proj_b

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q, b_q = _unpack_weight_bias(src.query)
                w_k, b_k = _unpack_weight_bias(src.key)
                w_v, b_v = _unpack_weight_bias(src.value)

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q, b_q = _unpack_weight_bias(src.q_proj)
                w_k, b_k = _unpack_weight_bias(src.k_proj)
                w_v, b_v = _unpack_weight_bias(src.v_proj)
                if hasattr(src, "o_proj"):
                    w_o, b_o = _unpack_weight_bias(src.o_proj)

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q, b_q = _unpack_weight_bias(src.q_lin)
                w_k, b_k = _unpack_weight_bias(src.k_lin)
                w_v, b_v = _unpack_weight_bias(src.v_lin)
                if hasattr(src, "out_lin"):
                    w_o, b_o = _unpack_weight_bias(src.out_lin)

            # Pattern 5: UniversalAdapter fallback
            if w_q is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    w_q, b_q = adapter.get_weight_bias(src, "q")
                    w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_k is None:
                        w_k, b_k = adapter.get_weight_bias(src, "k")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # Strict validation
            missing = []
            if w_q is None:
                missing.append("Q projection")
            if w_v is None:
                missing.append("V projection")
            if w_o is None and not self.is_nested_encoder:
                missing.append("Out projection")

            if missing:
                logger.warning(
                    f"SVDMamba: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
                )
                return

            if w_k is None:
                w_k = w_q

            # Copy into attention path
            self.q_proj.decompose_from_weight(
                w_q[: self.all_head_size, : self.d_model],
                b_q[: self.all_head_size] if b_q is not None else None,
            )
            self.k_proj.decompose_from_weight(
                w_k[: self.kv_head_size, : self.d_model],
                b_k[: self.kv_head_size] if b_k is not None else None,
            )
            self.v_proj.decompose_from_weight(
                w_v[: self.kv_head_size, : self.d_model],
                b_v[: self.kv_head_size] if b_v is not None else None,
            )

            # Assemble in_proj weight: [V, Q, V, Q] repeated to 2 * inner_dim
            if w_v.shape[0] < self.inner_dim:
                n_rep = math.ceil(self.inner_dim / w_v.shape[0])
                v_chunk = w_v.repeat(n_rep, 1)[: self.inner_dim, : self.d_model]
            else:
                v_chunk = w_v[: self.inner_dim, : self.d_model]

            if w_q.shape[0] < self.inner_dim:
                n_rep = math.ceil(self.inner_dim / w_q.shape[0])
                q_chunk = w_q.repeat(n_rep, 1)[: self.inner_dim, : self.d_model]
            else:
                q_chunk = w_q[: self.inner_dim, : self.d_model]

            in_w = torch.cat([v_chunk, q_chunk], dim=0)[: 2 * self.inner_dim, : self.d_model]

            in_b = None
            if b_v is not None and b_q is not None:
                if b_v.shape[0] < self.inner_dim:
                    n_rep = math.ceil(self.inner_dim / b_v.shape[0])
                    b_v_chunk = b_v.repeat(n_rep)[: self.inner_dim]
                else:
                    b_v_chunk = b_v[: self.inner_dim]

                if b_q.shape[0] < self.inner_dim:
                    n_rep = math.ceil(self.inner_dim / b_q.shape[0])
                    b_q_chunk = b_q.repeat(n_rep)[: self.inner_dim]
                else:
                    b_q_chunk = b_q[: self.inner_dim]

                in_b = torch.cat([b_v_chunk, b_q_chunk], dim=0)[: 2 * self.inner_dim]

            self.in_proj.decompose_from_weight(in_w, in_b)

        # Out proj (goes to attn_out_proj)
        if w_o is not None:
            out_w = w_o[: self.d_model, : self.all_head_size]
            if out_w.shape[1] < self.all_head_size:
                n_rep = math.ceil(self.all_head_size / out_w.shape[1])
                out_w = out_w.repeat(1, n_rep)[:, : self.all_head_size]
            self.attn_out_proj.decompose_from_weight(
                out_w, b_o[: self.d_model] if b_o is not None else None
            )

            # Verify Frobenius norms
            if float(torch.linalg.norm(self.in_proj.weight.float())) < 1e-4:
                logger.warning("SVDMamba: in_proj near-zero norm")
            if float(torch.linalg.norm(self.q_proj.weight.float())) < 1e-4:
                logger.warning("SVDMamba: q_proj near-zero norm")

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        if encoder_hidden_states is not None:
            return (hidden_states, None)

        while isinstance(hidden_states, (tuple, list)):
            hidden_states = hidden_states[0]

        if not hidden_states.is_contiguous():
            hidden_states = hidden_states.contiguous()

        batch, seqlen, _ = hidden_states.shape

        mask_bool = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                mask_bool = attention_mask.bool()
            elif attention_mask.dim() == 4:
                if attention_mask.shape[-2] == attention_mask.shape[-1]:
                    mask_bool = (attention_mask[:, 0, -1, :] > -10000.0).bool()
                else:
                    mask_bool = (attention_mask[:, 0, 0, :] > -10000.0).bool()
            elif attention_mask.dim() == 3:
                mask_bool = (
                    (attention_mask[:, 0, :] > -10000.0).bool()
                    if attention_mask.shape[1] > 1
                    else attention_mask.squeeze(1).bool()
                )

        if mask_bool is None:
            mask_bool = torch.ones(batch, seqlen, device=hidden_states.device, dtype=torch.bool)

        # ---------------------------------------------------------
        # 1. Attention Path (Multi-Head Attention with RoPE and Causal Masking)
        # ---------------------------------------------------------
        Q = (
            self.q_proj(hidden_states)
            .view(batch, seqlen, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            self.k_proj(hidden_states)
            .view(batch, seqlen, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V_attn = (
            self.v_proj(hidden_states)
            .view(batch, seqlen, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        pos_emb = (
            position_embeddings
            if position_embeddings is not None
            else kwargs.get("position_embeddings")
        )
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V_attn = repeat_kv(V_attn, self.kv_repeat)

        scale = 1.0 / math.sqrt(max(1, self.head_dim))
        attn_out = None
        if mask_bool 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_attn.half() if V_attn.dtype == torch.float32 else V_attn
                _is_causal = getattr(self, "is_decoder", False) and seqlen > 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,
                )
                attn_out = (
                    _sdpa_out.to(hidden_states.dtype)
                    .transpose(1, 2)
                    .contiguous()
                    .view(batch, seqlen, self.all_head_size)
                )
                attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
                attn_out = self.attn_out_proj(attn_out)
            except Exception:
                attn_out = None

        if attn_out is None:
            scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

            if self.is_decoder:
                causal_mask = torch.tril(
                    torch.ones(seqlen, seqlen, device=hidden_states.device, dtype=torch.bool)
                )
                scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4)

            if mask_bool is not None:
                scores = scores.masked_fill(~mask_bool.unsqueeze(1).unsqueeze(2), -1e4)

            probs = F.softmax(scores, dim=-1).to(V_attn.dtype)
            probs = torch.nan_to_num(probs, nan=0.0)
            attn_out = (
                torch.matmul(probs, V_attn)
                .transpose(1, 2)
                .contiguous()
                .view(batch, seqlen, self.all_head_size)
            )
            attn_out = self.attn_out_proj(attn_out)

        # ---------------------------------------------------------
        # 2. Selective SSM Path (Bounded FP32 associative update with min=1.0 denominator clamp)
        # ---------------------------------------------------------
        xz = self.in_proj(hidden_states)  # [B, L, 2*inner]
        x, z = xz.chunk(2, dim=-1)  # each [B, L, inner]

        x_t = x.transpose(1, 2).contiguous()  # [B, inner, L]
        x_pad = F.pad(x_t, (self.d_conv - 1, 0))
        x_conv = self.conv1d(x_pad)[..., :seqlen]
        x_conv = x_conv.transpose(1, 2).contiguous()  # [B, L, inner]
        x_conv = F.silu(x_conv)

        x_proj = self.x_proj(x_conv)  # [B, L, dt_rank+2*d_state]
        delta = x_proj[:, :, : self.dt_rank]
        B_ssm = x_proj[:, :, self.dt_rank : self.dt_rank + self.d_state]
        C_ssm = x_proj[:, :, self.dt_rank + self.d_state :]

        delta = self.dt_proj(delta)
        delta = F.softplus(delta + self.dt_bias).clamp(min=1e-4, max=1.0)

        A = -torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0)).to(dtype=x_conv.dtype)

        log_alpha = (delta.unsqueeze(-1) * A.view(1, 1, 1, -1)).float()
        u = (delta.unsqueeze(-1) * B_ssm.unsqueeze(2) * x_conv.unsqueeze(-1)).float()
        u_den = (delta.unsqueeze(-1) * B_ssm.unsqueeze(2).abs()).float()

        chunk_size = 64
        num_chunks = (seqlen + chunk_size - 1) // chunk_size
        h_running = x_conv.new_zeros(batch, self.inner_dim, self.d_state, dtype=torch.float32)
        h_den_running = x_conv.new_zeros(batch, self.inner_dim, self.d_state, dtype=torch.float32)
        y_chunks = []

        for c in range(num_chunks):
            s_idx = c * chunk_size
            e_idx = min(seqlen, s_idx + chunk_size)

            u_c = u[:, s_idx:e_idx]
            u_den_c = u_den[:, s_idx:e_idx]
            log_a_c = log_alpha[:, s_idx:e_idx]
            lambda_c = torch.clamp(torch.cumsum(log_a_c, dim=1), min=-25.0, max=0.0)

            decay_c = torch.exp(lambda_c)
            decay_inv_c = torch.exp(-lambda_c).clamp(max=1e10)

            h_intra = decay_c * torch.cumsum(decay_inv_c * u_c, dim=1)
            h_inter = decay_c * h_running.unsqueeze(1)
            h_chunk = torch.clamp(h_intra + h_inter, min=-1e4, max=1e4)

            h_den_intra = decay_c * torch.cumsum(decay_inv_c * u_den_c, dim=1)
            h_den_inter = decay_c * h_den_running.unsqueeze(1)
            h_den_chunk = torch.clamp(h_den_intra + h_den_inter, min=0.0, max=1e4)

            # Bounded denominator clamp to prevent explosive activations
            denom_c = (
                (C_ssm[:, s_idx:e_idx].abs().unsqueeze(2).float() * h_den_chunk)
                .sum(dim=-1)
                .clamp(min=1e-4)
            )
            num_c = (C_ssm[:, s_idx:e_idx].unsqueeze(2).float() * h_chunk).sum(dim=-1)
            y_c = (num_c / denom_c).to(x_conv.dtype)
            y_chunks.append(y_c)
            h_running = h_chunk[:, -1]
            h_den_running = h_den_chunk[:, -1]

        y = torch.cat(y_chunks, dim=1)
        y = y + self.D * x_conv
        ssm_out = y * F.silu(z)

        # Global SVD mixing
        global_scores = hidden_states @ self.global_V
        global_mix = self.global_out(global_scores)
        ssm_out = ssm_out + global_mix

        ssm_proj = self.out_proj(ssm_out)

        # ---------------------------------------------------------
        # 3. Dynamic Aggregation with Negative-Biased Gating
        # ---------------------------------------------------------
        gamma = torch.sigmoid(self.gating(hidden_states))
        Y = gamma * ssm_proj + (1.0 - gamma) * attn_out
        Y = torch.nan_to_num(Y, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            out = Y
        else:
            out = Y

        if mask_bool is not None:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        if out.shape[-1] != hidden_states.shape[-1]:
            if out.shape[-1] < hidden_states.shape[-1]:
                out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
            else:
                out = out[..., : hidden_states.shape[-1]]

        return (out, probs if output_attentions else None)

SVD-Mamba module with SVD-factorized projections, causal convolution, and selective SSM.

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.

Instance variables

prop k_lin : SVDLinear

Expand source code

@property
def k_lin(self) -> SVDLinear:
    return self.k_proj

prop key : SVDLinear

Expand source code

@property
def key(self) -> SVDLinear:
    return self.k_proj

prop q_lin : SVDLinear

Expand source code

@property
def q_lin(self) -> SVDLinear:
    return self.q_proj

prop query : SVDLinear

Expand source code

@property
def query(self) -> SVDLinear:
    return self.q_proj

prop v_lin : SVDLinear

Expand source code

@property
def v_lin(self) -> SVDLinear:
    return self.v_proj

prop value : SVDLinear

Expand source code

@property
def value(self) -> SVDLinear:
    return self.v_proj

Methods

def copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Check nested encoder container
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o, b_o = _unpack_weight_bias(old_attn.output.dense)
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
            if c_attn_w is not None:
                if (
                    c_attn_w.dim() == 2
                    and c_attn_w.shape[0] == self.d_model
                    and c_attn_w.shape[1] == 3 * self.d_model
                ):
                    c_attn_w = c_attn_w.t().contiguous()
                split_size = c_attn_w.size(0) // 3
                w_q = c_attn_w[:split_size, :]
                w_k = c_attn_w[split_size : 2 * split_size, :]
                w_v = c_attn_w[2 * split_size :, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:split_size]
                    b_k = c_attn_b[split_size : 2 * split_size]
                    b_v = c_attn_b[2 * split_size :]

            c_proj_w, c_proj_b = _unpack_weight_bias(src.c_proj)
            if c_proj_w is not None:
                if (
                    c_proj_w.dim() == 2
                    and c_proj_w.shape[0] != self.d_model
                    and c_proj_w.shape[1] == self.d_model
                ):
                    c_proj_w = c_proj_w.t().contiguous()
                w_o = c_proj_w
                b_o = c_proj_b

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q, b_q = _unpack_weight_bias(src.query)
            w_k, b_k = _unpack_weight_bias(src.key)
            w_v, b_v = _unpack_weight_bias(src.value)

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q, b_q = _unpack_weight_bias(src.q_proj)
            w_k, b_k = _unpack_weight_bias(src.k_proj)
            w_v, b_v = _unpack_weight_bias(src.v_proj)
            if hasattr(src, "o_proj"):
                w_o, b_o = _unpack_weight_bias(src.o_proj)

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q, b_q = _unpack_weight_bias(src.q_lin)
            w_k, b_k = _unpack_weight_bias(src.k_lin)
            w_v, b_v = _unpack_weight_bias(src.v_lin)
            if hasattr(src, "out_lin"):
                w_o, b_o = _unpack_weight_bias(src.out_lin)

        # Pattern 5: UniversalAdapter fallback
        if w_q is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                w_q, b_q = adapter.get_weight_bias(src, "q")
                w_v, b_v = adapter.get_weight_bias(src, "v")
                if w_k is None:
                    w_k, b_k = adapter.get_weight_bias(src, "k")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(src, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # Strict validation
        missing = []
        if w_q is None:
            missing.append("Q projection")
        if w_v is None:
            missing.append("V projection")
        if w_o is None and not self.is_nested_encoder:
            missing.append("Out projection")

        if missing:
            logger.warning(
                f"SVDMamba: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
            )
            return

        if w_k is None:
            w_k = w_q

        # Copy into attention path
        self.q_proj.decompose_from_weight(
            w_q[: self.all_head_size, : self.d_model],
            b_q[: self.all_head_size] if b_q is not None else None,
        )
        self.k_proj.decompose_from_weight(
            w_k[: self.kv_head_size, : self.d_model],
            b_k[: self.kv_head_size] if b_k is not None else None,
        )
        self.v_proj.decompose_from_weight(
            w_v[: self.kv_head_size, : self.d_model],
            b_v[: self.kv_head_size] if b_v is not None else None,
        )

        # Assemble in_proj weight: [V, Q, V, Q] repeated to 2 * inner_dim
        if w_v.shape[0] < self.inner_dim:
            n_rep = math.ceil(self.inner_dim / w_v.shape[0])
            v_chunk = w_v.repeat(n_rep, 1)[: self.inner_dim, : self.d_model]
        else:
            v_chunk = w_v[: self.inner_dim, : self.d_model]

        if w_q.shape[0] < self.inner_dim:
            n_rep = math.ceil(self.inner_dim / w_q.shape[0])
            q_chunk = w_q.repeat(n_rep, 1)[: self.inner_dim, : self.d_model]
        else:
            q_chunk = w_q[: self.inner_dim, : self.d_model]

        in_w = torch.cat([v_chunk, q_chunk], dim=0)[: 2 * self.inner_dim, : self.d_model]

        in_b = None
        if b_v is not None and b_q is not None:
            if b_v.shape[0] < self.inner_dim:
                n_rep = math.ceil(self.inner_dim / b_v.shape[0])
                b_v_chunk = b_v.repeat(n_rep)[: self.inner_dim]
            else:
                b_v_chunk = b_v[: self.inner_dim]

            if b_q.shape[0] < self.inner_dim:
                n_rep = math.ceil(self.inner_dim / b_q.shape[0])
                b_q_chunk = b_q.repeat(n_rep)[: self.inner_dim]
            else:
                b_q_chunk = b_q[: self.inner_dim]

            in_b = torch.cat([b_v_chunk, b_q_chunk], dim=0)[: 2 * self.inner_dim]

        self.in_proj.decompose_from_weight(in_w, in_b)

    # Out proj (goes to attn_out_proj)
    if w_o is not None:
        out_w = w_o[: self.d_model, : self.all_head_size]
        if out_w.shape[1] < self.all_head_size:
            n_rep = math.ceil(self.all_head_size / out_w.shape[1])
            out_w = out_w.repeat(1, n_rep)[:, : self.all_head_size]
        self.attn_out_proj.decompose_from_weight(
            out_w, b_o[: self.d_model] if b_o is not None else None
        )

        # Verify Frobenius norms
        if float(torch.linalg.norm(self.in_proj.weight.float())) < 1e-4:
            logger.warning("SVDMamba: in_proj near-zero norm")
        if float(torch.linalg.norm(self.q_proj.weight.float())) < 1e-4:
            logger.warning("SVDMamba: q_proj near-zero norm")

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class SVDMambaSub (d_model: int = 768,
config: Any | None = None,
rank: int | None = None,
is_decoder: bool = False,
**kwargs: Any)

Expand source code

class SVDMambaSub(nn.Module):
    is_custom_attention: bool = False

    def __init__(
        self,
        d_model: int = 768,
        config: Any | None = None,
        rank: int | None = None,
        is_decoder: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__()
        if d_model is None and config is not None:
            d_model = getattr(
                config,
                "hidden_size",
                getattr(config, "d_model", getattr(config, "dim", 768)),
            )
        d_model = int(d_model or 768)
        self.d_model = d_model
        self.ssm = SimpleSSM(d_model=d_model, config=config, is_decoder=is_decoder)
        self.in_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        # JAMBA-FIX-20: introspect SSM signature once
        self._ssm_supports_state = self._ssm_accepts_past_state()

    def _ssm_accepts_past_state(self) -> bool:
        import inspect

        try:
            sig = inspect.signature(self.ssm.forward)
            return "past_state" in sig.parameters or "state" in sig.parameters
        except (ValueError, TypeError):
            return False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        past_state: Any | None = None,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, Any | None]:
        x_proj = self.in_proj(hidden_states)

        ssm_result: Any = None
        present_state: Any | None = None

        if self._ssm_supports_state:
            try:
                ssm_result = self.ssm(x_proj, attention_mask=attention_mask, past_state=past_state)
            except TypeError:
                # Some SSMs call the kwarg `state` instead of `past_state`
                try:
                    ssm_result = self.ssm(x_proj, attention_mask=attention_mask, state=past_state)
                except Exception:
                    ssm_result = self.ssm(x_proj, attention_mask=attention_mask)
        else:
            ssm_result = self.ssm(x_proj, attention_mask=attention_mask)

        # Unpack (out, state) vs out
        if isinstance(ssm_result, tuple):
            ssm_out = ssm_result[0]
            if len(ssm_result) > 1:
                present_state = ssm_result[1]
        else:
            ssm_out = ssm_result

        out = self.out_proj(ssm_out)
        return out, present_state

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

Methods

def forward(self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
past_state: Any | None = None,
**kwargs: Any) ‑> tuple[torch.Tensor, Any | None]

Expand source code

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    past_state: Any | None = None,
    **kwargs: Any,
) -> tuple[torch.Tensor, Any | None]:
    x_proj = self.in_proj(hidden_states)

    ssm_result: Any = None
    present_state: Any | None = None

    if self._ssm_supports_state:
        try:
            ssm_result = self.ssm(x_proj, attention_mask=attention_mask, past_state=past_state)
        except TypeError:
            # Some SSMs call the kwarg `state` instead of `past_state`
            try:
                ssm_result = self.ssm(x_proj, attention_mask=attention_mask, state=past_state)
            except Exception:
                ssm_result = self.ssm(x_proj, attention_mask=attention_mask)
    else:
        ssm_result = self.ssm(x_proj, attention_mask=attention_mask)

    # Unpack (out, state) vs out
    if isinstance(ssm_result, tuple):
        ssm_out = ssm_result[0]
        if len(ssm_result) > 1:
            present_state = ssm_result[1]
    else:
        ssm_out = ssm_result

    out = self.out_proj(ssm_out)
    return out, present_state

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class SVDProjectionAttention (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
dropout: float = 0.1,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("svd_proj", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_proj_attn", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_projection_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("svd_projection", supports_causal=True, supports_gqa=True)
class SVDProjectionAttention(BaseAttention):
    r"""Multi-Head Attention using Factorized SVD Linear Projections.
    Hardened for production use across diverse transformer architectures.
    """

    supports_gqa: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "low_rank"

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        dropout: float = 0.1,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        # 1. Dimensions & Architecture
        self.hidden_size: int = kwargs.get(
            "hidden_size",
            _cfg_get(
                config,
                "hidden_size",
                _cfg_get(config, "dim", _cfg_get(config, "d_model", 768)),
            ),
        )
        self.num_heads: int = kwargs.get(
            "num_heads",
            kwargs.get(
                "num_attention_heads",
                _cfg_get(
                    config,
                    "num_attention_heads",
                    _cfg_get(config, "num_heads", _cfg_get(config, "n_heads", 12)),
                ),
            ),
        )
        self.head_dim: int = kwargs.get(
            "head_dim",
            _cfg_get(
                config,
                "head_dim",
                self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
            ),
        )

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads

        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat: int = self.num_heads // self.num_kv_heads
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.q_head_size: int = self.num_heads * self.head_dim

        # 2. Causal decoder flag detection
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(_cfg_get(config, "model_type", "") or "").lower()
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            self.is_decoder = (
                bool(_cfg_get(config, "is_decoder", False))
                or any(_mt.startswith(p) for p in _dec_pfx)
                or (old_attn is not None and hasattr(old_attn, "c_attn"))
                or bool(kwargs.get("is_causal", False))
            )

        self.is_nested_encoder: bool = False
        comp_config = _cfg_get(config, "compression_config", kwargs.get("comp_config"))
        self.rank: int = (
            rank
            if rank is not None
            else _cfg_get(
                comp_config,
                "rank",
                _cfg_get(config, "svd_rank", _cfg_get(config, "rank", 32)),
            )
        )
        self.lora_rank = _cfg_get(
            config,
            "lora_rank",
            _cfg_get(comp_config, "lora_rank", lora_rank if lora_rank is not None else 16),
        )
        self.lora_alpha = _cfg_get(
            config, "lora_alpha", _cfg_get(comp_config, "lora_alpha", lora_alpha)
        )

        # 3. Factorized SVD linear projections
        self.q_proj = SVDProjectionLinear(
            self.hidden_size,
            self.q_head_size,
            self.rank,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.k_proj = SVDProjectionLinear(
            self.hidden_size,
            self.kv_head_size,
            self.rank,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.v_proj = SVDProjectionLinear(
            self.hidden_size,
            self.kv_head_size,
            self.rank,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.out_proj = SVDProjectionLinear(
            self.q_head_size,
            self.hidden_size,
            self.rank,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = _cfg_get(
            config,
            "dropout",
            _cfg_get(
                config,
                "hidden_dropout_prob",
                _cfg_get(config, "attention_probs_dropout_prob", dropout),
            ),
        )
        self.dropout: nn.Dropout = nn.Dropout(dropout_p)
        self.triton_kwargs = {}

        # 4. Copy weights if old_attn passed
        if old_attn is not None:
            self.copy_weights(old_attn)
            try:
                src_device = next(old_attn.parameters()).device
                self.to(src_device)
            except StopIteration:
                pass

    @property
    def query(self) -> SVDProjectionLinear:
        return self.q_proj

    @property
    def key(self) -> SVDProjectionLinear:
        return self.k_proj

    @property
    def value(self) -> SVDProjectionLinear:
        return self.v_proj

    @property
    def q_lin(self) -> SVDProjectionLinear:
        return self.q_proj

    @property
    def k_lin(self) -> SVDProjectionLinear:
        return self.k_proj

    @property
    def v_lin(self) -> SVDProjectionLinear:
        return self.v_proj

    @property
    def out_lin(self) -> SVDProjectionLinear:
        return self.out_proj

    def _decompose_and_set(
        self,
        target: SVDProjectionLinear,
        weight: torch.Tensor,
        bias: torch.Tensor | None,
    ) -> None:
        """Decomposes full weight matrix into SVD factors and copies bias."""
        with torch.no_grad():
            W = weight.float()
            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
            r = min(target.rank, U.size(1), Vt.size(0))
            sqrt_s = torch.sqrt(torch.clamp(S[:r], min=1e-12))
            target.U.data.zero_()
            target.V.data.zero_()
            target.U.data[:, :r] = (U[:, :r] * sqrt_s.unsqueeze(0)).to(
                device=target.U.device, dtype=target.U.dtype
            )
            target.V.data[:r, :] = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(
                device=target.V.device, dtype=target.V.dtype
            )
            if target.bias is not None and bias is not None:
                target.bias.data.copy_(bias.to(device=target.bias.device, dtype=target.bias.dtype))

    def copy_weights(self, old_attn: nn.Module) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Detect nested encoder container or inner self-attention
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
                if c_attn_w is not None:
                    q_dim = self.q_head_size
                    k_dim = self.kv_head_size
                    v_dim = self.kv_head_size
                    if c_attn_w.size(0) >= q_dim + k_dim + v_dim:
                        w_q = c_attn_w[:q_dim, :]
                        w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                        w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
                        if c_attn_b is not None:
                            b_q = c_attn_b[:q_dim]
                            b_k = c_attn_b[q_dim : q_dim + k_dim]
                            b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]
                    else:
                        split_size = c_attn_w.size(0) // 3
                        w_q = c_attn_w[:split_size, :]
                        w_k = c_attn_w[split_size : 2 * split_size, :]
                        w_v = c_attn_w[2 * split_size : 3 * split_size, :]
                        if c_attn_b is not None:
                            b_q = c_attn_b[:split_size]
                            b_k = c_attn_b[split_size : 2 * split_size]
                            b_v = c_attn_b[2 * split_size : 3 * split_size]

                w_o, b_o = _unpack_weight_bias(src.c_proj)

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q, b_q = _unpack_weight_bias(src.query)
                w_k, b_k = _unpack_weight_bias(src.key)
                w_v, b_v = _unpack_weight_bias(src.value)

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q, b_q = _unpack_weight_bias(src.q_proj)
                w_k, b_k = _unpack_weight_bias(src.k_proj)
                w_v, b_v = _unpack_weight_bias(src.v_proj)
                if hasattr(src, "o_proj"):
                    w_o, b_o = _unpack_weight_bias(src.o_proj)

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q, b_q = _unpack_weight_bias(src.q_lin)
                w_k, b_k = _unpack_weight_bias(src.k_lin)
                w_v, b_v = _unpack_weight_bias(src.v_lin)
                if hasattr(src, "out_lin"):
                    w_o, b_o = _unpack_weight_bias(src.out_lin)

            # Fallback: UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    w_q, b_q = UniversalAdapter.get_weight_bias(old_attn, "q")
                    w_k, b_k = UniversalAdapter.get_weight_bias(old_attn, "k")
                    w_v, b_v = UniversalAdapter.get_weight_bias(old_attn, "v")
                    if w_o is None:
                        w_o, b_o = UniversalAdapter.get_weight_bias(old_attn, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # Verification of essential projections
            missing = []
            if w_q is None:
                missing.append("query/q_proj")
            if w_k is None:
                missing.append("key/k_proj")
            if w_v is None:
                missing.append("value/v_proj")
            if missing:
                logger.warning(f"Failed to extract required projections: {missing}")
                return

            # Frobenius norm checks
            for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    logger.warning(f"Copied weight {name} has near-zero Frobenius norm (< 1e-4)")

            # Decompose and set Q
            self._decompose_and_set(self.q_proj, w_q, b_q)

            # Handle GQA pooling/repetition for K and V
            def _align_kv(
                w: torch.Tensor, b: torch.Tensor | None
            ) -> tuple[torch.Tensor, torch.Tensor | None]:
                if w.shape[0] == self.kv_head_size:
                    return w, b
                elif w.shape[0] > self.kv_head_size and w.shape[0] % self.kv_head_size == 0:
                    rep = w.shape[0] // self.kv_head_size
                    w_pooled = (
                        w.view(self.num_kv_heads, rep, self.head_dim, -1)
                        .mean(dim=1)
                        .view(self.kv_head_size, -1)
                    )
                    b_pooled = (
                        b.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                        if b is not None
                        else None
                    )
                    return w_pooled, b_pooled
                elif w.shape[0] < self.kv_head_size and self.kv_head_size % w.shape[0] == 0:
                    rep = self.kv_head_size // w.shape[0]
                    w_rep = w.repeat(rep, 1)
                    b_rep = b.repeat(rep) if b is not None else None
                    return w_rep, b_rep
                else:
                    new_w = torch.zeros(
                        self.kv_head_size, w.shape[1], device=w.device, dtype=w.dtype
                    )
                    new_w[
                        : min(self.kv_head_size, w.shape[0]),
                        : min(self.hidden_size, w.shape[1]),
                    ].copy_(
                        w[
                            : min(self.kv_head_size, w.shape[0]),
                            : min(self.hidden_size, w.shape[1]),
                        ]
                    )
                    new_b = None
                    if b is not None:
                        new_b = torch.zeros(self.kv_head_size, device=b.device, dtype=b.dtype)
                        new_b[: min(self.kv_head_size, b.shape[0])].copy_(
                            b[: min(self.kv_head_size, b.shape[0])]
                        )
                    return new_w, new_b

            w_k_aligned, b_k_aligned = _align_kv(w_k, b_k)
            self._decompose_and_set(self.k_proj, w_k_aligned, b_k_aligned)

            w_v_aligned, b_v_aligned = _align_kv(w_v, b_v)
            self._decompose_and_set(self.v_proj, w_v_aligned, b_v_aligned)

            # Decompose and set out_proj if not nested encoder
            if w_o is not None and not self.is_nested_encoder:
                if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                    logger.warning("Copied out_proj weight has near-zero Frobenius norm (< 1e-4)")
                self._decompose_and_set(self.out_proj, w_o, b_o)

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        kv_len = kv_input.shape[1]

        # 1. Project Q, K, V
        Q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K = (
            self.k_proj(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V = (
            self.v_proj(kv_input)
            .view(B, kv_len, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. RoPE support
        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
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # 3. GQA expansion
        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        # 4. Attention scores in FP32
        scale = 1.0 / math.sqrt(self.head_dim)
        # Memory-efficient: use SDPA to avoid OOM on large seqs; fallback to explicit if needed

        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 = (
                getattr(self, "is_decoder", False) and not is_cross and N > 1
                if "is_cross" in dir()
                else getattr(self, "is_decoder", False) and N > 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, self.all_head_size)
            )

            _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

            return (
                self.out_proj(_sdpa_out)
                if not getattr(self, "is_nested_encoder", False)
                else _sdpa_out
            ), None

        except Exception:
            pass

        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        # 5. Masking
        if self.is_decoder and not is_cross and N > 1:
            causal_mask = (
                torch.triu(
                    torch.full(
                        (N, kv_len),
                        float("-inf"),
                        device=scores.device,
                        dtype=torch.float32,
                    ),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        attn_weights = self.dropout(attn_weights)

        # 6. Safe head_mask application
        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            attn_weights = attn_weights * head_mask

        # 7. Context computation & Head Merge
        context = torch.matmul(attn_weights, V).transpose(1, 2).contiguous().view(B, N, -1)
        valid_mask = get_token_valid_mask(eff_mask, B, N, hidden_states.device)
        context = context * valid_mask.unsqueeze(-1).to(context.dtype)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        # 8. Output projection (bypassed if nested encoder)
        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)

        if output.shape[-1] != hidden_states.shape[-1]:
            if output.shape[-1] < hidden_states.shape[-1]:
                output = F.pad(output, (0, hidden_states.shape[-1] - output.shape[-1]))
            else:
                output = output[..., : hidden_states.shape[-1]]

        return output, (attn_weights if output_attentions else None)

Multi-Head Attention using Factorized SVD Linear Projections. Hardened for production use across diverse transformer architectures.

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_gqa : bool

The type of the None singleton.

var triton_kernel : str

The type of the None singleton.

Instance variables

prop k_lin : SVDProjectionLinear

Expand source code

@property
def k_lin(self) -> SVDProjectionLinear:
    return self.k_proj

prop key : SVDProjectionLinear

Expand source code

@property
def key(self) -> SVDProjectionLinear:
    return self.k_proj

prop out_lin : SVDProjectionLinear

Expand source code

@property
def out_lin(self) -> SVDProjectionLinear:
    return self.out_proj

prop q_lin : SVDProjectionLinear

Expand source code

@property
def q_lin(self) -> SVDProjectionLinear:
    return self.q_proj

prop query : SVDProjectionLinear

Expand source code

@property
def query(self) -> SVDProjectionLinear:
    return self.q_proj

prop v_lin : SVDProjectionLinear

Expand source code

@property
def v_lin(self) -> SVDProjectionLinear:
    return self.v_proj

prop value : SVDProjectionLinear

Expand source code

@property
def value(self) -> SVDProjectionLinear:
    return self.v_proj

Methods

def copy_weights(self, old_attn: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Detect nested encoder container or inner self-attention
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight
                b_o = old_attn.output.dense.bias
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w, c_attn_b = _unpack_weight_bias(src.c_attn)
            if c_attn_w is not None:
                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size
                if c_attn_w.size(0) >= q_dim + k_dim + v_dim:
                    w_q = c_attn_w[:q_dim, :]
                    w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                    w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
                    if c_attn_b is not None:
                        b_q = c_attn_b[:q_dim]
                        b_k = c_attn_b[q_dim : q_dim + k_dim]
                        b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]
                else:
                    split_size = c_attn_w.size(0) // 3
                    w_q = c_attn_w[:split_size, :]
                    w_k = c_attn_w[split_size : 2 * split_size, :]
                    w_v = c_attn_w[2 * split_size : 3 * split_size, :]
                    if c_attn_b is not None:
                        b_q = c_attn_b[:split_size]
                        b_k = c_attn_b[split_size : 2 * split_size]
                        b_v = c_attn_b[2 * split_size : 3 * split_size]

            w_o, b_o = _unpack_weight_bias(src.c_proj)

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q, b_q = _unpack_weight_bias(src.query)
            w_k, b_k = _unpack_weight_bias(src.key)
            w_v, b_v = _unpack_weight_bias(src.value)

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q, b_q = _unpack_weight_bias(src.q_proj)
            w_k, b_k = _unpack_weight_bias(src.k_proj)
            w_v, b_v = _unpack_weight_bias(src.v_proj)
            if hasattr(src, "o_proj"):
                w_o, b_o = _unpack_weight_bias(src.o_proj)

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q, b_q = _unpack_weight_bias(src.q_lin)
            w_k, b_k = _unpack_weight_bias(src.k_lin)
            w_v, b_v = _unpack_weight_bias(src.v_lin)
            if hasattr(src, "out_lin"):
                w_o, b_o = _unpack_weight_bias(src.out_lin)

        # Fallback: UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                w_q, b_q = UniversalAdapter.get_weight_bias(old_attn, "q")
                w_k, b_k = UniversalAdapter.get_weight_bias(old_attn, "k")
                w_v, b_v = UniversalAdapter.get_weight_bias(old_attn, "v")
                if w_o is None:
                    w_o, b_o = UniversalAdapter.get_weight_bias(old_attn, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # Verification of essential projections
        missing = []
        if w_q is None:
            missing.append("query/q_proj")
        if w_k is None:
            missing.append("key/k_proj")
        if w_v is None:
            missing.append("value/v_proj")
        if missing:
            logger.warning(f"Failed to extract required projections: {missing}")
            return

        # Frobenius norm checks
        for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]:
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                logger.warning(f"Copied weight {name} has near-zero Frobenius norm (< 1e-4)")

        # Decompose and set Q
        self._decompose_and_set(self.q_proj, w_q, b_q)

        # Handle GQA pooling/repetition for K and V
        def _align_kv(
            w: torch.Tensor, b: torch.Tensor | None
        ) -> tuple[torch.Tensor, torch.Tensor | None]:
            if w.shape[0] == self.kv_head_size:
                return w, b
            elif w.shape[0] > self.kv_head_size and w.shape[0] % self.kv_head_size == 0:
                rep = w.shape[0] // self.kv_head_size
                w_pooled = (
                    w.view(self.num_kv_heads, rep, self.head_dim, -1)
                    .mean(dim=1)
                    .view(self.kv_head_size, -1)
                )
                b_pooled = (
                    b.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                    if b is not None
                    else None
                )
                return w_pooled, b_pooled
            elif w.shape[0] < self.kv_head_size and self.kv_head_size % w.shape[0] == 0:
                rep = self.kv_head_size // w.shape[0]
                w_rep = w.repeat(rep, 1)
                b_rep = b.repeat(rep) if b is not None else None
                return w_rep, b_rep
            else:
                new_w = torch.zeros(
                    self.kv_head_size, w.shape[1], device=w.device, dtype=w.dtype
                )
                new_w[
                    : min(self.kv_head_size, w.shape[0]),
                    : min(self.hidden_size, w.shape[1]),
                ].copy_(
                    w[
                        : min(self.kv_head_size, w.shape[0]),
                        : min(self.hidden_size, w.shape[1]),
                    ]
                )
                new_b = None
                if b is not None:
                    new_b = torch.zeros(self.kv_head_size, device=b.device, dtype=b.dtype)
                    new_b[: min(self.kv_head_size, b.shape[0])].copy_(
                        b[: min(self.kv_head_size, b.shape[0])]
                    )
                return new_w, new_b

        w_k_aligned, b_k_aligned = _align_kv(w_k, b_k)
        self._decompose_and_set(self.k_proj, w_k_aligned, b_k_aligned)

        w_v_aligned, b_v_aligned = _align_kv(w_v, b_v)
        self._decompose_and_set(self.v_proj, w_v_aligned, b_v_aligned)

        # Decompose and set out_proj if not nested encoder
        if w_o is not None and not self.is_nested_encoder:
            if torch.norm(w_o.float(), p="fro").item() < 1e-4:
                logger.warning("Copied out_proj weight has near-zero Frobenius norm (< 1e-4)")
            self._decompose_and_set(self.out_proj, w_o, b_o)

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class SVDProjectionLinear (in_features: int,
out_features: int,
rank: int,
bias: bool = True,
lora_rank: int = 16,
lora_alpha: float = 32.0,
freeze_base: bool = True)

Expand source code

class SVDProjectionLinear(nn.Module):
    is_custom_attention: bool = True
    r"""Low-Rank SVD Linear Layer with Trainable LoRA Residual Adapters.

    Forward computes:
    $$y = x (U V)^\top + b + \frac{\alpha}{r_{\text{lora}}} x A^\top B^\top$$
    where $U \in \mathbb{R}^{d_{\text{out}} \times r}$, $V \in \mathbb{R}^{r \times d_{\text{in}}}$
    are frozen SVD singular subspaces, and $A, B$ are trainable LoRA adapters.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        freeze_base: bool = True,
    ) -> None:
        super().__init__()
        self.in_features: int = in_features
        self.out_features: int = out_features
        self.rank: int = rank
        self.lora_rank: int = lora_rank
        self.lora_alpha: float = lora_alpha
        self.scaling: float = lora_alpha / max(1, lora_rank)

        self.U: nn.Parameter = nn.Parameter(torch.empty(out_features, rank))
        self.V: nn.Parameter = nn.Parameter(torch.empty(rank, in_features))
        if bias:
            self.bias: nn.Parameter | None = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

        if lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty(lora_rank, in_features))
            self.lora_B = nn.Parameter(torch.empty(out_features, lora_rank))
            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)

        if freeze_base:
            self.U.requires_grad = False
            self.V.requires_grad = False
            if self.bias is not None:
                self.bias.requires_grad = False

    def reset_parameters(self) -> None:
        """Initialize SVD parameters with Kaiming uniform distribution."""
        nn.init.kaiming_uniform_(self.U, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.V, a=math.sqrt(5))
        if self.bias is not None:
            fan_in = self.in_features
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)

    @property
    def weight(self) -> torch.Tensor:
        """Reconstructed effective dense weight matrix."""
        return torch.matmul(self.U, self.V)

    triton_kernel: str = "low_rank"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Base factorized SVD projection
        out = triton_dispatch(self.triton_kernel, x, self.U, self.V, bias=self.bias)

        # Trainable LoRA residual branch
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scaling
            out = out + lora_out

        return out

    def set_svd_factors(self, U: torch.Tensor, V: torch.Tensor) -> None:
        """Set precomputed U and V factors."""
        with torch.no_grad():
            r = min(self.rank, U.size(1), V.size(0))
            self.U.data.zero_()
            self.V.data.zero_()
            self.U.data[:, :r].copy_(U[:, :r].to(device=self.U.device, dtype=self.U.dtype))
            self.V.data[:r, :].copy_(V[:r, :].to(device=self.V.device, dtype=self.V.dtype))

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

Low-Rank SVD Linear Layer with Trainable LoRA Residual Adapters.

Forward computes: $$y = x (U V)^\top + b + \frac{\alpha}{r_{\text{lora}}} x A^\top B^\top$$ where $U \in \mathbb{R}^{d_{\text{out}} \times r}$, $V \in \mathbb{R}^{r \times d_{\text{in}}}$ are frozen SVD singular subspaces, and $A, B$ are trainable LoRA adapters.

var triton_kernel : str

The type of the None singleton.

Instance variables

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Reconstructed effective dense weight matrix."""
    return torch.matmul(self.U, self.V)

Reconstructed effective dense weight matrix.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # Base factorized SVD projection
    out = triton_dispatch(self.triton_kernel, x, self.U, self.V, bias=self.bias)

    # Trainable LoRA residual branch
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scaling
        out = out + lora_out

    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    """Initialize SVD parameters with Kaiming uniform distribution."""
    nn.init.kaiming_uniform_(self.U, a=math.sqrt(5))
    nn.init.kaiming_uniform_(self.V, a=math.sqrt(5))
    if self.bias is not None:
        fan_in = self.in_features
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        nn.init.uniform_(self.bias, -bound, bound)

Initialize SVD parameters with Kaiming uniform distribution.

def set_svd_factors(self, U: torch.Tensor, V: torch.Tensor) ‑> None

Expand source code

def set_svd_factors(self, U: torch.Tensor, V: torch.Tensor) -> None:
    """Set precomputed U and V factors."""
    with torch.no_grad():
        r = min(self.rank, U.size(1), V.size(0))
        self.U.data.zero_()
        self.V.data.zero_()
        self.U.data[:, :r].copy_(U[:, :r].to(device=self.U.device, dtype=self.U.dtype))
        self.V.data[:r, :].copy_(V[:r, :].to(device=self.V.device, dtype=self.V.dtype))

Set precomputed U and V factors.

class SVDS6StateSpace (d_model: int,
rank: int = 16,
bottleneck_dim: int = 64,
dtype: torch.dtype | None = None,
device: torch.device | None = None)

Expand source code

class SVDS6StateSpace(nn.Module):
    """
    Per-token low-rank SVD state transition.

        U_t, Σ_t, V_t, B_t, C_t = Bottleneck(x_t)
        h_t = U_t · diag(Σ_t) · V_t · h_{t-1} + B_t · x_t
        y_t = C_t · h_t

    Contractive by construction (MSF-FIX-03):
      ||U_t||_2 ≤ 1, ||V_t||_2 ≤ 1, Σ_t ∈ (0, 1)   ⇒  ||UΣV||_2 ≤ 1.
    """

    def __init__(
        self,
        d_model: int,
        rank: int = 16,
        bottleneck_dim: int = 64,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ):
        super().__init__()
        self.d_model = d_model
        self.rank = max(1, int(rank))
        self.bottleneck_dim = max(self.rank, int(bottleneck_dim))

        fk = {"device": device, "dtype": dtype}

        # MSF-FIX-10: bottleneck to keep params small
        self.down = nn.Linear(d_model, self.bottleneck_dim, bias=False, **fk)

        # Outputs from the bottleneck:
        #   - U:   d_model * rank         (per-token basis for columns)
        #   - V:   rank * d_model         (per-token basis for rows)
        #   - sigma: rank                 (per-token singular values)
        #   - B:   d_model                (input gate)
        #   - C:   d_model                (output gate)
        self.up_u = nn.Linear(self.bottleneck_dim, d_model * self.rank, bias=False, **fk)
        self.up_v = nn.Linear(self.bottleneck_dim, self.rank * d_model, bias=False, **fk)
        self.up_sigma = nn.Linear(self.bottleneck_dim, self.rank, bias=True, **fk)
        self.up_b = nn.Linear(self.bottleneck_dim, d_model, bias=True, **fk)
        self.up_c = nn.Linear(self.bottleneck_dim, d_model, bias=True, **fk)

        # Sigma bias: initial sigmoid ≈ 0.85 (stable contraction)
        with torch.no_grad():
            nn.init.zeros_(self.up_sigma.weight)
            nn.init.constant_(self.up_sigma.bias, 1.75)  # sigmoid(1.75) ≈ 0.85

        # C bias: identity-like start
        with torch.no_grad():
            if self.up_c.bias is not None:
                nn.init.zeros_(self.up_c.bias)

    def forward(self, x: torch.Tensor):
        B, L, D = x.shape
        r = self.rank

        z = self.down(x)  # (B, L, bottleneck)

        U = self.up_u(z).view(B, L, D, r)
        V = self.up_v(z).view(B, L, r, D)
        sigma = torch.sigmoid(self.up_sigma(z))  # (B, L, r) ∈ (0,1)
        B_seq = self.up_b(z)  # (B, L, D)
        C_seq = self.up_c(z)  # (B, L, D)

        # MSF-FIX-03: enforce unit spectral norms on U, V
        U = F.normalize(U, p=2, dim=-2, eps=1e-8)
        V = F.normalize(V, p=2, dim=-2, eps=1e-8)

        return U, sigma, V, B_seq, C_seq

Per-token low-rank SVD state transition.

U_t, Σ_t, V_t, B_t, C_t = Bottleneck(x_t)
h_t = U_t · diag(Σ_t) · V_t · h_{t-1} + B_t · x_t
y_t = C_t · h_t

Contractive by construction (MSF-FIX-03): ||U_t||_2 ≤ 1, ||V_t||_2 ≤ 1, Σ_t ∈ (0, 1) ⇒ ||UΣV||_2 ≤ 1.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor) ‑> Callable[..., typing.Any]

Expand source code

def forward(self, x: torch.Tensor):
    B, L, D = x.shape
    r = self.rank

    z = self.down(x)  # (B, L, bottleneck)

    U = self.up_u(z).view(B, L, D, r)
    V = self.up_v(z).view(B, L, r, D)
    sigma = torch.sigmoid(self.up_sigma(z))  # (B, L, r) ∈ (0,1)
    B_seq = self.up_b(z)  # (B, L, D)
    C_seq = self.up_c(z)  # (B, L, D)

    # MSF-FIX-03: enforce unit spectral norms on U, V
    U = F.normalize(U, p=2, dim=-2, eps=1e-8)
    V = F.normalize(V, p=2, dim=-2, eps=1e-8)

    return U, sigma, V, B_seq, C_seq

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class SambaMoEBlock (config: Any | None = None,
d_model: int | None = None,
n_heads: int | None = None,
n_experts: int = 8,
window: int = 128,
rank: int | None = 64,
pretrained_attn: torch.nn.modules.module.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("samba_moe", supports_causal=True, supports_gqa=True)
class SambaMoEBlock(BaseAttention):
    r"""SambaMoE: SAMBA-M3 + REAP-MoE Hybrid Architecture with pretrained weight copying."""

    supports_gqa: bool = True
    supports_causal: bool = True

    def __init__(
        self,
        config: Any | None = None,
        d_model: int | None = None,
        n_heads: int | None = None,
        n_experts: int = 8,
        window: int = 128,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        hidden_size = kwargs.get("d_model") or kwargs.get("hidden_size") or d_model
        if pretrained_attn is not None:
            cand = getattr(
                pretrained_attn,
                "hidden_size",
                getattr(
                    pretrained_attn,
                    "embed_dim",
                    getattr(pretrained_attn, "all_head_size", None),
                ),
            )
            if (
                cand is None
                and hasattr(pretrained_attn, "q_proj")
                and hasattr(pretrained_attn.q_proj, "in_features")
            ):
                cand = pretrained_attn.q_proj.in_features
            elif (
                cand is None
                and hasattr(pretrained_attn, "q_lin")
                and hasattr(pretrained_attn.q_lin, "in_features")
            ):
                cand = pretrained_attn.q_lin.in_features
            if cand is not None:
                hidden_size = cand
        if hidden_size is None and config is not None:
            hidden_size = _cfg_get(config, "hidden_size", _cfg_get(config, "d_model", None))
        hidden_size = int(hidden_size or 768)

        if config is not None:
            n_heads = _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "n_heads", n_heads or 12),
            )
            n_experts = _cfg_get(config, "num_experts", _cfg_get(config, "n_experts", n_experts))
            window = _cfg_get(config, "window_size", window)
        else:
            n_heads = n_heads or 12

        self.d_model = hidden_size
        self.n_heads = int(n_heads)
        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.swa_mamba = SAMBAM3Block(
            config=config,
            d_model=self.d_model,
            n_heads=self.n_heads,
            window=window,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            is_decoder=is_decoder,
            **kwargs,
        )

        if REAPMoE is not None:
            try:
                self.moe = REAPMoE(self.d_model, n_experts=n_experts)
            except Exception as e:
                logger.warning("REAPMoE construction failed: %s, using MinimalMoEFallback", e)
                self.moe = MinimalMoEFallback(self.d_model, n_experts=min(4, n_experts))
        else:
            self.moe = MinimalMoEFallback(self.d_model, n_experts=min(4, n_experts))

        self.norm = nn.LayerNorm(self.d_model)
        self.aux_loss: torch.Tensor = torch.tensor(0.0)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    @property
    def is_nested_encoder(self) -> bool:
        return getattr(self.swa_mamba, "is_nested_encoder", False)

    @is_nested_encoder.setter
    def is_nested_encoder(self, val: bool) -> None:
        if hasattr(self, "swa_mamba"):
            self.swa_mamba.is_nested_encoder = val

    @property
    def query(self) -> nn.Module:
        return getattr(self.swa_mamba, "q_lin", None)

    @property
    def key(self) -> nn.Module:
        return getattr(self.swa_mamba, "k_lin", None)

    @property
    def value(self) -> nn.Module:
        return getattr(self.swa_mamba, "v_lin", None)

    def copy_weights(self, pretrained_attn: nn.Module) -> None:
        """Propagate weight transfer to internal SAMBA-M3 and SSM components."""
        src_cls_name = getattr(pretrained_attn, "__class__", type(pretrained_attn)).__name__
        has_out_proj = any(
            hasattr(pretrained_attn, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
            )
        )
        if (
            src_cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (src_cls_name.endswith("SelfAttention") and not has_out_proj)
            or (hasattr(pretrained_attn, "query") and not has_out_proj)
            or (hasattr(pretrained_attn, "self") and hasattr(pretrained_attn, "output"))
        ):
            self.is_nested_encoder = True
            if hasattr(self, "swa_mamba"):
                self.swa_mamba.is_nested_encoder = True

        if hasattr(self.swa_mamba, "copy_weights"):
            self.swa_mamba.copy_weights(pretrained_attn)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        if encoder_hidden_states is not None:
            raise NotImplementedError("SambaMoEBlock does not support cross-attention")

        x, attn_w = self.swa_mamba(
            hidden_states,
            attention_mask=attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            **kwargs,
        )

        router_in = getattr(self.moe, "d_model", self.d_model)
        if hasattr(self.moe, "router") and hasattr(self.moe.router, "in_features"):
            router_in = self.moe.router.in_features

        if x.shape[-1] != router_in:
            x_moe = (
                F.pad(x, (0, router_in - x.shape[-1]))
                if x.shape[-1] < router_in
                else x[..., :router_in]
            )
        else:
            x_moe = x

        moe_out, _ = self.moe(x_moe)
        self.aux_loss = getattr(self.moe, "aux_loss", torch.tensor(0.0, device=x.device))

        if moe_out.shape[-1] != x.shape[-1]:
            moe_out = (
                F.pad(moe_out, (0, x.shape[-1] - moe_out.shape[-1]))
                if moe_out.shape[-1] < x.shape[-1]
                else moe_out[..., : x.shape[-1]]
            )

        target_dtype = hidden_states.dtype
        x = x.to(target_dtype)
        moe_out = moe_out.to(target_dtype)

        if self.is_nested_encoder:
            # Nested encoder: bypass outer LayerNorm (self.norm) and double residual
            out = x + moe_out
        else:
            out = self.norm(x + moe_out).to(target_dtype)

        if out.shape[-1] != hidden_states.shape[-1]:
            out = (
                F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
                if out.shape[-1] < hidden_states.shape[-1]
                else out[..., : hidden_states.shape[-1]]
            )
        return (out.to(target_dtype), attn_w)

SambaMoE: SAMBA-M3 + REAP-MoE Hybrid Architecture with pretrained weight copying.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• BaseAttention
• torch.nn.modules.module.Module

Class variables

var supports_causal : bool

The type of the None singleton.

var supports_gqa : bool

The type of the None singleton.

Instance variables

prop is_nested_encoder : bool

Expand source code

@property
def is_nested_encoder(self) -> bool:
    return getattr(self.swa_mamba, "is_nested_encoder", False)

prop key : torch.nn.modules.module.Module

Expand source code

@property
def key(self) -> nn.Module:
    return getattr(self.swa_mamba, "k_lin", None)

prop query : torch.nn.modules.module.Module

Expand source code

@property
def query(self) -> nn.Module:
    return getattr(self.swa_mamba, "q_lin", None)

prop value : torch.nn.modules.module.Module

Expand source code

@property
def value(self) -> nn.Module:
    return getattr(self.swa_mamba, "v_lin", None)

Methods

def copy_weights(self, pretrained_attn: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, pretrained_attn: nn.Module) -> None:
    """Propagate weight transfer to internal SAMBA-M3 and SSM components."""
    src_cls_name = getattr(pretrained_attn, "__class__", type(pretrained_attn)).__name__
    has_out_proj = any(
        hasattr(pretrained_attn, attr)
        for attr in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
        )
    )
    if (
        src_cls_name in ("BertSelfAttention", "RobertaSelfAttention")
        or (src_cls_name.endswith("SelfAttention") and not has_out_proj)
        or (hasattr(pretrained_attn, "query") and not has_out_proj)
        or (hasattr(pretrained_attn, "self") and hasattr(pretrained_attn, "output"))
    ):
        self.is_nested_encoder = True
        if hasattr(self, "swa_mamba"):
            self.swa_mamba.is_nested_encoder = True

    if hasattr(self.swa_mamba, "copy_weights"):
        self.swa_mamba.copy_weights(pretrained_attn)

Propagate weight transfer to internal SAMBA-M3 and SSM components.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class SeerGating (head_dim: int,
num_heads: int,
gate_rank: int,
top_k: int,
temperature: float = 1.0,
**kwargs)

Expand source code

class SeerGating(nn.Module):
    def __init__(
        self,
        head_dim: int,
        num_heads: int,
        gate_rank: int,
        top_k: int,
        temperature: float = 1.0,
        **kwargs,
    ) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.top_k = top_k
        self.temperature = temperature

        self.q_pool = nn.Sequential(
            nn.Linear(head_dim, gate_rank),
            nn.GELU(),
            nn.Linear(gate_rank, gate_rank),
        )
        self.k_pool = nn.Sequential(
            nn.Linear(head_dim, gate_rank),
            nn.GELU(),
            nn.Linear(gate_rank, gate_rank),
        )

    def forward(
        self, Q: torch.Tensor, K: torch.Tensor, n_blocks: int, training: bool = True
    ) -> torch.Tensor:
        B, H, N, D = Q.shape
        gr = self.q_pool[0].out_features

        Q_pooled = Q.mean(dim=2)
        K_pooled = K.mean(dim=2)

        q_gate = self.q_pool(Q_pooled).view(B, H, 1, gr)
        k_gate = self.k_pool(K_pooled).view(B, H, gr, 1)

        gate_score = torch.sigmoid(
            torch.matmul(q_gate, k_gate).squeeze(-1).squeeze(-1) / self.temperature
        )

        keep_per_head = torch.clamp((gate_score * n_blocks).long(), min=1, max=n_blocks)

        mask = torch.zeros(B, H, n_blocks, n_blocks, device=Q.device, dtype=torch.bool)
        diag_idx = torch.arange(n_blocks, device=Q.device)
        mask[:, :, diag_idx, diag_idx] = True
        mask[:, :, :, 0] = True

        for b in range(B):
            for h in range(H):
                k_keep = keep_per_head[b, h].item()
                if k_keep > 1 and n_blocks > 2:
                    extra_needed = min(k_keep - 1, n_blocks - 1)
                    if extra_needed > 0:
                        mask[b, h, :, 1 : 1 + extra_needed] = True
        return mask

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, Q: torch.Tensor, K: torch.Tensor, n_blocks: int, training: bool = True) ‑> torch.Tensor

Expand source code

def forward(
    self, Q: torch.Tensor, K: torch.Tensor, n_blocks: int, training: bool = True
) -> torch.Tensor:
    B, H, N, D = Q.shape
    gr = self.q_pool[0].out_features

    Q_pooled = Q.mean(dim=2)
    K_pooled = K.mean(dim=2)

    q_gate = self.q_pool(Q_pooled).view(B, H, 1, gr)
    k_gate = self.k_pool(K_pooled).view(B, H, gr, 1)

    gate_score = torch.sigmoid(
        torch.matmul(q_gate, k_gate).squeeze(-1).squeeze(-1) / self.temperature
    )

    keep_per_head = torch.clamp((gate_score * n_blocks).long(), min=1, max=n_blocks)

    mask = torch.zeros(B, H, n_blocks, n_blocks, device=Q.device, dtype=torch.bool)
    diag_idx = torch.arange(n_blocks, device=Q.device)
    mask[:, :, diag_idx, diag_idx] = True
    mask[:, :, :, 0] = True

    for b in range(B):
        for h in range(H):
            k_keep = keep_per_head[b, h].item()
            if k_keep > 1 and n_blocks > 2:
                extra_needed = min(k_keep - 1, n_blocks - 1)
                if extra_needed > 0:
                    mask[b, h, :, 1 : 1 + extra_needed] = True
    return mask

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class SelectiveMamba (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
dim: int | None = None,
d_state: int = 16,
d_conv: int = 4,
expand: int = 1,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("selective_mamba", supports_causal=True, supports_bidirectional=True)
@ATTENTION_REGISTRY.register(
    "upi_mamba", supports_causal=True, supports_bidirectional=True
)  # DEPRECATED alias
@ATTENTION_REGISTRY.register(
    "upi_stuffed_mamba", supports_causal=True, supports_bidirectional=True
)  # DEPRECATED alias
class SelectiveMamba(BaseAttention):
    r"""
    Selective Mamba SSM block.

    Implements a Mamba-style selective state space model with
    input-dependent B, C, and dt parameters:
        h_t = Ā_t · h_{t-1} + B̄_t · x_t
        y_t = C_t · h_t + D · x_t
    where Ā_t = exp(dt_t · A).

    The "selectivity" comes from x_proj producing B_t and C_t
    from the input.

    NOTE: The previous name "UPI-Stuffed" has been removed
    because none of the claimed features (Uncertainty Penalty,
    Interpolation, Stuffing) were actually implemented.

    Registered as "upi_mamba" (legacy alias) and
    "selective_mamba" (canonical).
    """

    supports_causal: bool = True
    supports_bidirectional: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        dim: int | None = None,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 1,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        src_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")

        if dim is None and d_model is None:
            if config is not None and not isinstance(config, int):
                dim = _cfg_get(
                    config,
                    "hidden_size",
                    _cfg_get(config, "dim", kwargs.get("hidden_size", 768)),
                )
            else:
                dim = kwargs.get("hidden_size", 768)
        else:
            dim = dim or d_model

        if hasattr(dim, "hidden_size"):
            dim = dim.hidden_size

        self.dim: int = int(dim)
        self.d_state: int = d_state
        self.d_conv: int = d_conv
        self.expand: int = expand
        self.inner_dim: int = self.dim * expand
        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))

        if is_decoder is not None:
            self.is_decoder: bool = is_decoder
        elif config is not None and not isinstance(config, int):
            self.is_decoder = bool(getattr(config, "is_decoder", False))
        else:
            self.is_decoder = False

        self.is_nested_encoder: bool = False

        self.in_proj = LoRALinear(
            self.dim,
            self.inner_dim * 2,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.conv1d = nn.Conv1d(
            self.inner_dim,
            self.inner_dim,
            kernel_size=d_conv,
            groups=self.inner_dim,
            padding=d_conv - 1,
        )
        # A_log init: log([1, 2, ..., d_state]) repeated.
        # In our scan, we compute A = +exp(A_log) (positive),
        # then log_decay = -dt * A (negative), giving decay < 1.
        # This is mathematically equivalent to Mamba's
        # A = -exp(A_log), decay = exp(dt * A).
        A_init = torch.arange(1, d_state + 1, dtype=torch.float32).repeat(self.inner_dim, 1)
        self.A_log = nn.Parameter(torch.log(A_init))

        # Trainable skip connection D matching inner_dim
        self.D = nn.Parameter(torch.ones(self.inner_dim))

        # Mamba paper: dt_rank = ceil(d_model / 16), not d_inner
        self.dt_rank = math.ceil(self.dim / 16)

        # x_proj maps inner_dim → dt_rank + 2*d_state
        self.x_proj = nn.Linear(
            self.inner_dim,
            self.dt_rank + 2 * self.d_state,
            bias=False,
        )
        self.x_proj.weight.requires_grad = True

        # dt_proj maps dt_rank → inner_dim
        self.dt_proj = nn.Linear(self.dt_rank, self.inner_dim, bias=True)

        # Initialize dt_proj bias (Mamba paper Eq. 12)
        dt_min, dt_max = 0.001, 0.1
        dt_init = torch.exp(
            torch.rand(self.inner_dim) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)
        )
        dt_init_inv = dt_init + torch.log(-torch.expm1(-dt_init))
        with torch.no_grad():
            self.dt_proj.bias.copy_(dt_init_inv)
            # Weight init: uniform(-1/sqrt(dt_rank), 1/sqrt(dt_rank))
            bound = 1.0 / math.sqrt(self.dt_rank)
            nn.init.uniform_(self.dt_proj.weight, -bound, bound)
        self.dt_proj.weight.requires_grad = True
        self.dt_proj.bias.requires_grad = True

        self.out_proj = LoRALinear(
            self.inner_dim,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        if src_attn is not None:
            self.copy_weights(src_attn)

    @property
    def query(self) -> LoRALinear:
        return self.in_proj

    @property
    def key(self) -> LoRALinear:
        return self.in_proj

    @property
    def value(self) -> LoRALinear:
        return self.in_proj

    @property
    def q_lin(self) -> LoRALinear:
        return self.in_proj

    @property
    def k_lin(self) -> LoRALinear:
        return self.in_proj

    @property
    def v_lin(self) -> LoRALinear:
        return self.in_proj

    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, _b_q = None, None
            w_k, _b_k = None, None
            w_v, _b_v = None, None
            w_o, b_o = None, None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight.detach()
                    b_o = (
                        old_attn.output.dense.bias.detach()
                        if old_attn.output.dense.bias is not None
                        else None
                    )
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                (src.c_attn.bias.detach() if src.c_attn.bias is not None else None)
                if c_attn_w.dim() == 2 and (
                    c_attn_w.shape[0] == self.dim and c_attn_w.shape[1] != self.dim
                ):
                    c_attn_w = c_attn_w.t()

                split_q = min(self.dim, c_attn_w.shape[0] // 3)
                w_q = c_attn_w[:split_q, :]
                w_k = c_attn_w[split_q : 2 * split_q, :]
                w_v = c_attn_w[2 * split_q : 3 * split_q, :]

                out_w = src.c_proj.weight.detach()
                if out_w.dim() == 2 and (out_w.shape[0] != self.dim and out_w.shape[1] == self.dim):
                    out_w = out_w.t()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                src.value.bias.detach() if src.value.bias is not None else None

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                src.v_proj.bias.detach() if src.v_proj.bias is not None else None

                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # Pattern 4: DistilBERT
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                src.v_lin.bias.detach() if src.v_lin.bias is not None else None

                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # Fallback: UniversalAdapter extraction
            if w_q is None or w_v is None or (w_o is None and not self.is_nested_encoder):
                for proj in ["q", "k", "v", "o"]:
                    try:
                        w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
                        if w is not None:
                            if proj == "q" and w_q is None:
                                w_q, _b_q = w, b
                            elif proj == "k" and w_k is None:
                                w_k, _b_k = w, b
                            elif proj == "v" and w_v is None:
                                w_v, _b_v = w, b
                            elif proj == "o" and w_o is None:
                                w_o, b_o = w, b
                    except Exception as e:
                        logger.debug("UniversalAdapter extraction for %s failed: %s", proj, e)

            missing = []
            if w_q is None:
                missing.append("Q")
            if w_v is None:
                missing.append("V")
            if w_o is None and not self.is_nested_encoder:
                missing.append("O")
            if missing:
                raise RuntimeError(
                    f"SelectiveMamba: cannot extract {missing} from "
                    f"{old_attn.__class__.__name__}. Aborting to avoid "
                    f"silent model corruption."
                )

            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()
                return w

            def _adapt_to_inner(w: torch.Tensor, name: str) -> torch.Tensor:
                """
                Adapt weight from source head count to inner_dim.
                If source is smaller: tile and truncate.
                If source is larger: slice to fit.
                """
                src_out, src_in = w.shape
                if src_out == self.inner_dim:
                    return w

                if src_out > self.inner_dim:
                    # Source has more heads — take the first inner_dim
                    logger.debug(
                        "SelectiveMamba: slicing %s from %d to %d",
                        name,
                        src_out,
                        self.inner_dim,
                    )
                    return w[: self.inner_dim, :].contiguous()

                # Source is smaller — repeat to fill
                rep = math.ceil(self.inner_dim / src_out)
                logger.debug(
                    "SelectiveMamba: repeating %s from %d to %d (rep=%d)",
                    name,
                    src_out,
                    self.inner_dim,
                    rep,
                )
                return w.repeat(rep, 1)[: self.inner_dim, :].contiguous()

            w_v_or = _orient(w_v, self.dim)
            w_q_or = _orient(w_q, self.dim)
            w_v_or = _adapt_to_inner(w_v_or, "W_V")
            w_q_or = _adapt_to_inner(w_q_or, "W_Q")

            v_chunk = w_v_or[:, : self.dim]
            q_chunk = w_q_or[:, : self.dim]

            self.in_proj.weight.data[: self.inner_dim, : self.dim].copy_(v_chunk)
            self.in_proj.weight.data[self.inner_dim : 2 * self.inner_dim, : self.dim].copy_(q_chunk)

            if w_o is not None:
                w_o_or = _orient(w_o, self.inner_dim)
                min_o_out = min(self.dim, w_o_or.shape[0])
                min_o_in = min(self.dim, w_o_or.shape[1])
                self.out_proj.weight.data[:min_o_out, :min_o_in].copy_(
                    w_o_or[:min_o_out, :min_o_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))]
                    )

            # Frobenius norm checks
            for name, w in [
                ("in_proj", self.in_proj.weight),
                ("out_proj", self.out_proj.weight),
            ]:
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    logger.warning(f"SelectiveMamba: {name} has near-zero Frobenius norm (< 1e-4)")

            self.in_proj.freeze_base_weights()
            self.out_proj.freeze_base_weights()
            self.x_proj.weight.requires_grad = True
            self.dt_proj.weight.requires_grad = True
            self.dt_proj.bias.requires_grad = True

            # Preserve source LoRA adapters
            source_submodules = {
                "in": ["q_proj", "q_lin", "query", "W_q"],
                "out": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
            }
            for name in source_submodules["in"]:
                if hasattr(old_attn, name):
                    sub = getattr(old_attn, name)
                    if (
                        hasattr(sub, "lora_A")
                        and hasattr(sub, "lora_B")
                        and sub.lora_A is not None
                        and sub.lora_B is not None
                    ):
                        src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                        src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                        if hasattr(src_lA, "data") and self.in_proj.lora_A is not None:
                            r_min = min(self.in_proj.lora_A.shape[0], src_lA.shape[0])
                            in_min = min(self.in_proj.lora_A.shape[1], src_lA.shape[1])
                            self.in_proj.lora_A.data[:r_min, :in_min].copy_(
                                src_lA.data[:r_min, :in_min]
                            )
                            self.in_proj.lora_A.requires_grad = True
                        if hasattr(src_lB, "data") and self.in_proj.lora_B is not None:
                            out_min = min(self.in_proj.lora_B.shape[0], src_lB.shape[0])
                            r_min = min(self.in_proj.lora_B.shape[1], src_lB.shape[1])
                            self.in_proj.lora_B.data[:out_min, :r_min].copy_(
                                src_lB.data[:out_min, :r_min]
                            )
                            self.in_proj.lora_B.requires_grad = True
                        if hasattr(sub, "lora_scaling"):
                            self.in_proj.scaling = sub.lora_scaling
                        break

            for name in source_submodules["out"]:
                if hasattr(old_attn, name):
                    sub = getattr(old_attn, name)
                    if (
                        hasattr(sub, "lora_A")
                        and hasattr(sub, "lora_B")
                        and sub.lora_A is not None
                        and sub.lora_B is not None
                    ):
                        src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                        src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                        if hasattr(src_lA, "data") and self.out_proj.lora_A is not None:
                            r_min = min(self.out_proj.lora_A.shape[0], src_lA.shape[0])
                            in_min = min(self.out_proj.lora_A.shape[1], src_lA.shape[1])
                            self.out_proj.lora_A.data[:r_min, :in_min].copy_(
                                src_lA.data[:r_min, :in_min]
                            )
                            self.out_proj.lora_A.requires_grad = True
                        if hasattr(src_lB, "data") and self.out_proj.lora_B is not None:
                            out_min = min(self.out_proj.lora_B.shape[0], src_lB.shape[0])
                            r_min = min(self.out_proj.lora_B.shape[1], src_lB.shape[1])
                            self.out_proj.lora_B.data[:out_min, :r_min].copy_(
                                src_lB.data[:out_min, :r_min]
                            )
                            self.out_proj.lora_B.requires_grad = True
                        if hasattr(sub, "lora_scaling"):
                            self.out_proj.scaling = sub.lora_scaling
                        break

    def _selective_scan(
        self,
        x: torch.Tensor,
        dt: torch.Tensor,
        B_t: torch.Tensor,
        C_t: torch.Tensor,
        mask: torch.Tensor | None = None,
        reverse: bool = False,
    ) -> torch.Tensor:
        """Bounded chunked associative SSM scan in FP32 with O(chunk=64) memory."""
        if reverse:
            x = torch.flip(x, dims=[1])
            dt = torch.flip(dt, dims=[1])
            B_t = torch.flip(B_t, dims=[1])
            C_t = torch.flip(C_t, dims=[1])
            if mask is not None:
                mask = torch.flip(mask, dims=[1])

        batch, seqlen, _ = x.shape
        A = torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0))  # [inner_dim, d_state]

        # Path 1: Hardware-fused Triton selective scan on CUDA
        if HAS_TRITON_UPI and triton_selective_scan is not None and x.is_cuda and mask is None:
            try:
                return triton_selective_scan(
                    x.contiguous(),
                    dt.contiguous(),
                    -A.contiguous(),
                    B_t.contiguous(),
                    C_t.contiguous(),
                    None,
                    delta_softplus=False,
                ).to(x.dtype)
            except Exception as e:
                logger.debug("Triton selective scan failed, falling back to chunked PyTorch: %s", e)

        # In case of padding mask, zero out invalid token inputs
        if mask is not None:
            m = mask.unsqueeze(-1).float()
            x = x * m
            dt = dt * m

        chunk_size = 64
        num_chunks = (seqlen + chunk_size - 1) // chunk_size
        h_running = x.new_zeros(batch, self.inner_dim, self.d_state, dtype=torch.float32)
        y_chunks = []
        A_view = A.view(1, 1, self.inner_dim, self.d_state)

        for c in range(num_chunks):
            s_idx = c * chunk_size
            e_idx = min(seqlen, s_idx + chunk_size)

            dt_c = dt[:, s_idx:e_idx]
            x_c = x[:, s_idx:e_idx]
            B_c = B_t[:, s_idx:e_idx]
            C_c = C_t[:, s_idx:e_idx]

            # Compute u and log_a strictly for current 64-token chunk to avoid buffer overflow
            log_a_c = -(dt_c.unsqueeze(-1) * A_view).float()
            u_c = (dt_c.unsqueeze(-1) * B_c.unsqueeze(2) * x_c.unsqueeze(-1)).float()

            lambda_c = torch.clamp(torch.cumsum(log_a_c, dim=1), min=-25.0, max=0.0)

            decay_c = torch.exp(lambda_c)
            decay_inv_c = torch.exp(-lambda_c).clamp(max=1e10)

            h_intra = decay_c * torch.cumsum(decay_inv_c * u_c, dim=1)
            h_inter = decay_c * h_running.unsqueeze(1)

            h_chunk = torch.clamp(h_intra + h_inter, min=-1e4, max=1e4)
            y_c = (C_c.unsqueeze(2).float() * h_chunk).sum(dim=-1).to(x.dtype)
            y_chunks.append(y_c)
            h_running = h_chunk[:, -1]

        y = torch.cat(y_chunks, dim=1)
        y = (y + self.D.to(x.dtype) * x).to(x.dtype)

        if reverse:
            y = torch.flip(y, dims=[1])

        return y

    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]:
        if encoder_hidden_states is not None:
            raise NotImplementedError(
                "SelectiveMamba does not support cross-attention. "
                "This module replaces SELF-attention only. "
                "For cross-attention, use a different module."
            )

        B, L, D = hidden_states.shape
        mask = attention_mask if attention_mask is not None else kwargs.get("mask")

        if head_mask is not None:
            logger.warning(
                "SelectiveMamba does not use head_mask (SSM has no heads). Ignoring head_mask."
            )

        mask_bool = None
        if mask is not None:
            if mask.dim() == 2:
                # Distinguish between int/bool masks and float additive masks
                if mask.dtype in (
                    torch.bool,
                    torch.int8,
                    torch.int16,
                    torch.int32,
                    torch.int64,
                ):
                    mask_bool = mask.bool().to(hidden_states.device)
                else:
                    # Float mask: -inf = invalid, 0 = valid
                    mask_bool = (mask > -1e4).to(hidden_states.device)
            elif mask.dim() == 4:
                mask_bool = (mask[:, 0, 0, :] > -1e4).to(hidden_states.device)
            else:
                try:
                    mask_bool = get_token_valid_mask(
                        mask,
                        B,
                        L,
                        hidden_states.device,
                    )
                except Exception as e:
                    logger.debug("get_token_valid_mask failed: %s", e)
                    mask_bool = None

        target_dtype = hidden_states.dtype
        xz = self.in_proj(hidden_states)
        x_in, z = xz.chunk(2, dim=-1)

        conv_w = self.conv1d.weight.to(x_in.dtype)
        conv_b = self.conv1d.bias.to(x_in.dtype) if self.conv1d.bias is not None else None
        x_conv = (
            F.conv1d(
                x_in.transpose(1, 2),
                conv_w,
                conv_b,
                stride=self.conv1d.stride,
                padding=self.conv1d.padding,
                dilation=self.conv1d.dilation,
                groups=self.conv1d.groups,
            )[:, :, :L]
            .transpose(1, 2)
            .contiguous()
        )
        x_ssm = F.silu(x_conv)

        # Selective projections: dt_raw [dt_rank], B_t [d_state], C_t [d_state]
        x_proj_out = self.x_proj(x_ssm)
        dt_raw, B_t, C_t = torch.split(
            x_proj_out,
            [self.dt_rank, self.d_state, self.d_state],
            dim=-1,
        )
        dt = F.softplus(self.dt_proj(dt_raw))
        dt = torch.clamp(dt, min=1e-3, max=1.0)

        mt = str(getattr(getattr(self, "config", None), "model_type", "") or "").lower()
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_decoder = (
            self.is_decoder
            or getattr(getattr(self, "config", None), "is_decoder", False)
            or any(mt.startswith(p) for p in decoder_prefixes)
        )

        y_fwd = self._selective_scan(x_ssm, dt, B_t, C_t, mask=mask_bool, reverse=False)

        if is_decoder:
            y = y_fwd
        else:
            y_bwd = self._selective_scan(x_ssm, dt, B_t, C_t, mask=mask_bool, reverse=True)
            y = 0.5 * (y_fwd + y_bwd)

        gated = (y.to(target_dtype) * F.silu(z.to(target_dtype))).to(target_dtype)

        # Inner SSM output is inner_dim; out_proj brings it back to dim.
        # Always apply out_proj.
        out = self.out_proj(gated).to(target_dtype)

        if mask_bool is not None:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        return (out, None)

    def step_recurrent(
        self,
        x_t: torch.Tensor,
        state: dict[str, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """
        Single-step streaming inference.

        x_t: (B, 1, D)
        state: dict with keys 'h', 'conv_buf' (optional)
        Returns: (y_t, new_state)
        """
        if x_t.dim() == 2:
            x_t = x_t.unsqueeze(1)
        B, L, D = x_t.shape
        assert L == 1, "step_recurrent expects L=1"

        # Initialize state
        if state is None:
            h = torch.zeros(
                B,
                self.inner_dim,
                self.d_state,
                device=x_t.device,
                dtype=torch.float32,
            )
            conv_buf = torch.zeros(
                B,
                self.d_conv - 1,
                self.inner_dim,
                device=x_t.device,
                dtype=x_t.dtype,
            )
        else:
            h = state["h"]
            conv_buf = state.get("conv_buf")

        # in_proj
        xz = self.in_proj(x_t)  # (B, 1, 2*inner_dim)
        x_in, z = xz.chunk(2, dim=-1)

        # Conv1d with buffer
        x_conv_input = torch.cat([conv_buf, x_in], dim=1)  # (B, d_conv, inner_dim)
        x_conv = F.conv1d(
            x_conv_input.transpose(1, 2),
            self.conv1d.weight,
            self.conv1d.bias,
            groups=self.inner_dim,
        ).transpose(1, 2)  # (B, 1, inner_dim)
        new_conv_buf = x_conv_input[:, -(self.d_conv - 1) :]

        x_ssm = F.silu(x_conv)

        # Selective params
        x_proj_out = self.x_proj(x_ssm)
        dt_raw, B_t, C_t = torch.split(
            x_proj_out,
            [self.dt_rank, self.d_state, self.d_state],
            dim=-1,
        )
        dt = F.softplus(self.dt_proj(dt_raw)).clamp(1e-3, 1.0)

        # Single step scan
        A = torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0))
        dA = torch.exp(-dt.unsqueeze(-1) * A.view(1, 1, self.inner_dim, self.d_state))
        dBx = dt.unsqueeze(-1) * B_t.unsqueeze(2) * x_ssm.unsqueeze(-1)

        h = dA.squeeze(1) * h + dBx.squeeze(1)  # (B, inner_dim, d_state)
        y = (C_t.squeeze(1).unsqueeze(1) * h).sum(dim=-1)  # (B, inner_dim)
        y = y + self.D * x_ssm.squeeze(1)

        # Gating
        gated = y * F.silu(z.squeeze(1))  # (B, inner_dim)

        # Output projection
        out = self.out_proj(gated.unsqueeze(1))  # (B, 1, D)

        new_state = {"h": h, "conv_buf": new_conv_buf}
        return out, new_state

Selective Mamba SSM block.

Implements a Mamba-style selective state space model with input-dependent B, C, and dt parameters: h_t = Ā_t · h_{t-1} + B̄_t · x_t y_t = C_t · h_t + D · x_t where Ā_t = exp(dt_t · A).

The "selectivity" comes from x_proj producing B_t and C_t from the input.

NOTE: The previous name "UPI-Stuffed" has been removed because none of the claimed features (Uncertainty Penalty, Interpolation, Stuffing) were actually implemented.

Registered as "upi_mamba" (legacy alias) and "selective_mamba" (canonical).

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.

Instance variables

prop k_lin : LoRALinear

Expand source code

@property
def k_lin(self) -> LoRALinear:
    return self.in_proj

prop key : LoRALinear

Expand source code

@property
def key(self) -> LoRALinear:
    return self.in_proj

prop q_lin : LoRALinear

Expand source code

@property
def q_lin(self) -> LoRALinear:
    return self.in_proj

prop query : LoRALinear

Expand source code

@property
def query(self) -> LoRALinear:
    return self.in_proj

prop v_lin : LoRALinear

Expand source code

@property
def v_lin(self) -> LoRALinear:
    return self.in_proj

prop value : LoRALinear

Expand source code

@property
def value(self) -> LoRALinear:
    return self.in_proj

Methods

def copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, _b_q = None, None
        w_k, _b_k = None, None
        w_v, _b_v = None, None
        w_o, b_o = None, None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight.detach()
                b_o = (
                    old_attn.output.dense.bias.detach()
                    if old_attn.output.dense.bias is not None
                    else None
                )
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            (src.c_attn.bias.detach() if src.c_attn.bias is not None else None)
            if c_attn_w.dim() == 2 and (
                c_attn_w.shape[0] == self.dim and c_attn_w.shape[1] != self.dim
            ):
                c_attn_w = c_attn_w.t()

            split_q = min(self.dim, c_attn_w.shape[0] // 3)
            w_q = c_attn_w[:split_q, :]
            w_k = c_attn_w[split_q : 2 * split_q, :]
            w_v = c_attn_w[2 * split_q : 3 * split_q, :]

            out_w = src.c_proj.weight.detach()
            if out_w.dim() == 2 and (out_w.shape[0] != self.dim and out_w.shape[1] == self.dim):
                out_w = out_w.t()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            src.value.bias.detach() if src.value.bias is not None else None

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            src.v_proj.bias.detach() if src.v_proj.bias is not None else None

            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # Pattern 4: DistilBERT
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            src.v_lin.bias.detach() if src.v_lin.bias is not None else None

            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # Fallback: UniversalAdapter extraction
        if w_q is None or w_v is None or (w_o is None and not self.is_nested_encoder):
            for proj in ["q", "k", "v", "o"]:
                try:
                    w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
                    if w is not None:
                        if proj == "q" and w_q is None:
                            w_q, _b_q = w, b
                        elif proj == "k" and w_k is None:
                            w_k, _b_k = w, b
                        elif proj == "v" and w_v is None:
                            w_v, _b_v = w, b
                        elif proj == "o" and w_o is None:
                            w_o, b_o = w, b
                except Exception as e:
                    logger.debug("UniversalAdapter extraction for %s failed: %s", proj, e)

        missing = []
        if w_q is None:
            missing.append("Q")
        if w_v is None:
            missing.append("V")
        if w_o is None and not self.is_nested_encoder:
            missing.append("O")
        if missing:
            raise RuntimeError(
                f"SelectiveMamba: cannot extract {missing} from "
                f"{old_attn.__class__.__name__}. Aborting to avoid "
                f"silent model corruption."
            )

        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()
            return w

        def _adapt_to_inner(w: torch.Tensor, name: str) -> torch.Tensor:
            """
            Adapt weight from source head count to inner_dim.
            If source is smaller: tile and truncate.
            If source is larger: slice to fit.
            """
            src_out, src_in = w.shape
            if src_out == self.inner_dim:
                return w

            if src_out > self.inner_dim:
                # Source has more heads — take the first inner_dim
                logger.debug(
                    "SelectiveMamba: slicing %s from %d to %d",
                    name,
                    src_out,
                    self.inner_dim,
                )
                return w[: self.inner_dim, :].contiguous()

            # Source is smaller — repeat to fill
            rep = math.ceil(self.inner_dim / src_out)
            logger.debug(
                "SelectiveMamba: repeating %s from %d to %d (rep=%d)",
                name,
                src_out,
                self.inner_dim,
                rep,
            )
            return w.repeat(rep, 1)[: self.inner_dim, :].contiguous()

        w_v_or = _orient(w_v, self.dim)
        w_q_or = _orient(w_q, self.dim)
        w_v_or = _adapt_to_inner(w_v_or, "W_V")
        w_q_or = _adapt_to_inner(w_q_or, "W_Q")

        v_chunk = w_v_or[:, : self.dim]
        q_chunk = w_q_or[:, : self.dim]

        self.in_proj.weight.data[: self.inner_dim, : self.dim].copy_(v_chunk)
        self.in_proj.weight.data[self.inner_dim : 2 * self.inner_dim, : self.dim].copy_(q_chunk)

        if w_o is not None:
            w_o_or = _orient(w_o, self.inner_dim)
            min_o_out = min(self.dim, w_o_or.shape[0])
            min_o_in = min(self.dim, w_o_or.shape[1])
            self.out_proj.weight.data[:min_o_out, :min_o_in].copy_(
                w_o_or[:min_o_out, :min_o_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))]
                )

        # Frobenius norm checks
        for name, w in [
            ("in_proj", self.in_proj.weight),
            ("out_proj", self.out_proj.weight),
        ]:
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                logger.warning(f"SelectiveMamba: {name} has near-zero Frobenius norm (< 1e-4)")

        self.in_proj.freeze_base_weights()
        self.out_proj.freeze_base_weights()
        self.x_proj.weight.requires_grad = True
        self.dt_proj.weight.requires_grad = True
        self.dt_proj.bias.requires_grad = True

        # Preserve source LoRA adapters
        source_submodules = {
            "in": ["q_proj", "q_lin", "query", "W_q"],
            "out": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
        }
        for name in source_submodules["in"]:
            if hasattr(old_attn, name):
                sub = getattr(old_attn, name)
                if (
                    hasattr(sub, "lora_A")
                    and hasattr(sub, "lora_B")
                    and sub.lora_A is not None
                    and sub.lora_B is not None
                ):
                    src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                    src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                    if hasattr(src_lA, "data") and self.in_proj.lora_A is not None:
                        r_min = min(self.in_proj.lora_A.shape[0], src_lA.shape[0])
                        in_min = min(self.in_proj.lora_A.shape[1], src_lA.shape[1])
                        self.in_proj.lora_A.data[:r_min, :in_min].copy_(
                            src_lA.data[:r_min, :in_min]
                        )
                        self.in_proj.lora_A.requires_grad = True
                    if hasattr(src_lB, "data") and self.in_proj.lora_B is not None:
                        out_min = min(self.in_proj.lora_B.shape[0], src_lB.shape[0])
                        r_min = min(self.in_proj.lora_B.shape[1], src_lB.shape[1])
                        self.in_proj.lora_B.data[:out_min, :r_min].copy_(
                            src_lB.data[:out_min, :r_min]
                        )
                        self.in_proj.lora_B.requires_grad = True
                    if hasattr(sub, "lora_scaling"):
                        self.in_proj.scaling = sub.lora_scaling
                    break

        for name in source_submodules["out"]:
            if hasattr(old_attn, name):
                sub = getattr(old_attn, name)
                if (
                    hasattr(sub, "lora_A")
                    and hasattr(sub, "lora_B")
                    and sub.lora_A is not None
                    and sub.lora_B is not None
                ):
                    src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                    src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                    if hasattr(src_lA, "data") and self.out_proj.lora_A is not None:
                        r_min = min(self.out_proj.lora_A.shape[0], src_lA.shape[0])
                        in_min = min(self.out_proj.lora_A.shape[1], src_lA.shape[1])
                        self.out_proj.lora_A.data[:r_min, :in_min].copy_(
                            src_lA.data[:r_min, :in_min]
                        )
                        self.out_proj.lora_A.requires_grad = True
                    if hasattr(src_lB, "data") and self.out_proj.lora_B is not None:
                        out_min = min(self.out_proj.lora_B.shape[0], src_lB.shape[0])
                        r_min = min(self.out_proj.lora_B.shape[1], src_lB.shape[1])
                        self.out_proj.lora_B.data[:out_min, :r_min].copy_(
                            src_lB.data[:out_min, :r_min]
                        )
                        self.out_proj.lora_B.requires_grad = True
                    if hasattr(sub, "lora_scaling"):
                        self.out_proj.scaling = sub.lora_scaling
                    break

Robust weight copying from DistilBERT, RoBERTa, GPT-2 Conv1D, or Qwen/LLaMA.

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]

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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    if encoder_hidden_states is not None:
        raise NotImplementedError(
            "SelectiveMamba does not support cross-attention. "
            "This module replaces SELF-attention only. "
            "For cross-attention, use a different module."
        )

    B, L, D = hidden_states.shape
    mask = attention_mask if attention_mask is not None else kwargs.get("mask")

    if head_mask is not None:
        logger.warning(
            "SelectiveMamba does not use head_mask (SSM has no heads). Ignoring head_mask."
        )

    mask_bool = None
    if mask is not None:
        if mask.dim() == 2:
            # Distinguish between int/bool masks and float additive masks
            if mask.dtype in (
                torch.bool,
                torch.int8,
                torch.int16,
                torch.int32,
                torch.int64,
            ):
                mask_bool = mask.bool().to(hidden_states.device)
            else:
                # Float mask: -inf = invalid, 0 = valid
                mask_bool = (mask > -1e4).to(hidden_states.device)
        elif mask.dim() == 4:
            mask_bool = (mask[:, 0, 0, :] > -1e4).to(hidden_states.device)
        else:
            try:
                mask_bool = get_token_valid_mask(
                    mask,
                    B,
                    L,
                    hidden_states.device,
                )
            except Exception as e:
                logger.debug("get_token_valid_mask failed: %s", e)
                mask_bool = None

    target_dtype = hidden_states.dtype
    xz = self.in_proj(hidden_states)
    x_in, z = xz.chunk(2, dim=-1)

    conv_w = self.conv1d.weight.to(x_in.dtype)
    conv_b = self.conv1d.bias.to(x_in.dtype) if self.conv1d.bias is not None else None
    x_conv = (
        F.conv1d(
            x_in.transpose(1, 2),
            conv_w,
            conv_b,
            stride=self.conv1d.stride,
            padding=self.conv1d.padding,
            dilation=self.conv1d.dilation,
            groups=self.conv1d.groups,
        )[:, :, :L]
        .transpose(1, 2)
        .contiguous()
    )
    x_ssm = F.silu(x_conv)

    # Selective projections: dt_raw [dt_rank], B_t [d_state], C_t [d_state]
    x_proj_out = self.x_proj(x_ssm)
    dt_raw, B_t, C_t = torch.split(
        x_proj_out,
        [self.dt_rank, self.d_state, self.d_state],
        dim=-1,
    )
    dt = F.softplus(self.dt_proj(dt_raw))
    dt = torch.clamp(dt, min=1e-3, max=1.0)

    mt = str(getattr(getattr(self, "config", None), "model_type", "") or "").lower()
    decoder_prefixes = (
        "gpt",
        "qwen",
        "llama",
        "mistral",
        "gemma",
        "phi",
        "opt",
        "bloom",
        "falcon",
        "olmo",
        "mpt",
        "stablelm",
    )
    is_decoder = (
        self.is_decoder
        or getattr(getattr(self, "config", None), "is_decoder", False)
        or any(mt.startswith(p) for p in decoder_prefixes)
    )

    y_fwd = self._selective_scan(x_ssm, dt, B_t, C_t, mask=mask_bool, reverse=False)

    if is_decoder:
        y = y_fwd
    else:
        y_bwd = self._selective_scan(x_ssm, dt, B_t, C_t, mask=mask_bool, reverse=True)
        y = 0.5 * (y_fwd + y_bwd)

    gated = (y.to(target_dtype) * F.silu(z.to(target_dtype))).to(target_dtype)

    # Inner SSM output is inner_dim; out_proj brings it back to dim.
    # Always apply out_proj.
    out = self.out_proj(gated).to(target_dtype)

    if mask_bool is not None:
        out = out * mask_bool.unsqueeze(-1).to(out.dtype)

    return (out, None)

Must return (output, attention_weights_or_None).

def step_recurrent(self, x_t: torch.Tensor, state: dict[str, torch.Tensor] | None = None) ‑> tuple[torch.Tensor, dict[str, torch.Tensor]]

Expand source code

def step_recurrent(
    self,
    x_t: torch.Tensor,
    state: dict[str, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """
    Single-step streaming inference.

    x_t: (B, 1, D)
    state: dict with keys 'h', 'conv_buf' (optional)
    Returns: (y_t, new_state)
    """
    if x_t.dim() == 2:
        x_t = x_t.unsqueeze(1)
    B, L, D = x_t.shape
    assert L == 1, "step_recurrent expects L=1"

    # Initialize state
    if state is None:
        h = torch.zeros(
            B,
            self.inner_dim,
            self.d_state,
            device=x_t.device,
            dtype=torch.float32,
        )
        conv_buf = torch.zeros(
            B,
            self.d_conv - 1,
            self.inner_dim,
            device=x_t.device,
            dtype=x_t.dtype,
        )
    else:
        h = state["h"]
        conv_buf = state.get("conv_buf")

    # in_proj
    xz = self.in_proj(x_t)  # (B, 1, 2*inner_dim)
    x_in, z = xz.chunk(2, dim=-1)

    # Conv1d with buffer
    x_conv_input = torch.cat([conv_buf, x_in], dim=1)  # (B, d_conv, inner_dim)
    x_conv = F.conv1d(
        x_conv_input.transpose(1, 2),
        self.conv1d.weight,
        self.conv1d.bias,
        groups=self.inner_dim,
    ).transpose(1, 2)  # (B, 1, inner_dim)
    new_conv_buf = x_conv_input[:, -(self.d_conv - 1) :]

    x_ssm = F.silu(x_conv)

    # Selective params
    x_proj_out = self.x_proj(x_ssm)
    dt_raw, B_t, C_t = torch.split(
        x_proj_out,
        [self.dt_rank, self.d_state, self.d_state],
        dim=-1,
    )
    dt = F.softplus(self.dt_proj(dt_raw)).clamp(1e-3, 1.0)

    # Single step scan
    A = torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0))
    dA = torch.exp(-dt.unsqueeze(-1) * A.view(1, 1, self.inner_dim, self.d_state))
    dBx = dt.unsqueeze(-1) * B_t.unsqueeze(2) * x_ssm.unsqueeze(-1)

    h = dA.squeeze(1) * h + dBx.squeeze(1)  # (B, inner_dim, d_state)
    y = (C_t.squeeze(1).unsqueeze(1) * h).sum(dim=-1)  # (B, inner_dim)
    y = y + self.D * x_ssm.squeeze(1)

    # Gating
    gated = y * F.silu(z.squeeze(1))  # (B, inner_dim)

    # Output projection
    out = self.out_proj(gated.unsqueeze(1))  # (B, 1, D)

    new_state = {"h": h, "conv_buf": new_conv_buf}
    return out, new_state

Single-step streaming inference.

x_t: (B, 1, D) state: dict with keys 'h', 'conv_buf' (optional) Returns: (y_t, new_state)

class UPIMambaBlock (config: Any | None = None,
rank: int | None = None,
d_model: int | Any | None = None,
dim: int | None = None,
d_state: int = 16,
d_conv: int = 4,
expand: int = 1,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("selective_mamba", supports_causal=True, supports_bidirectional=True)
@ATTENTION_REGISTRY.register(
    "upi_mamba", supports_causal=True, supports_bidirectional=True
)  # DEPRECATED alias
@ATTENTION_REGISTRY.register(
    "upi_stuffed_mamba", supports_causal=True, supports_bidirectional=True
)  # DEPRECATED alias
class SelectiveMamba(BaseAttention):
    r"""
    Selective Mamba SSM block.

    Implements a Mamba-style selective state space model with
    input-dependent B, C, and dt parameters:
        h_t = Ā_t · h_{t-1} + B̄_t · x_t
        y_t = C_t · h_t + D · x_t
    where Ā_t = exp(dt_t · A).

    The "selectivity" comes from x_proj producing B_t and C_t
    from the input.

    NOTE: The previous name "UPI-Stuffed" has been removed
    because none of the claimed features (Uncertainty Penalty,
    Interpolation, Stuffing) were actually implemented.

    Registered as "upi_mamba" (legacy alias) and
    "selective_mamba" (canonical).
    """

    supports_causal: bool = True
    supports_bidirectional: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | Any | None = None,
        dim: int | None = None,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 1,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        src_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")

        if dim is None and d_model is None:
            if config is not None and not isinstance(config, int):
                dim = _cfg_get(
                    config,
                    "hidden_size",
                    _cfg_get(config, "dim", kwargs.get("hidden_size", 768)),
                )
            else:
                dim = kwargs.get("hidden_size", 768)
        else:
            dim = dim or d_model

        if hasattr(dim, "hidden_size"):
            dim = dim.hidden_size

        self.dim: int = int(dim)
        self.d_state: int = d_state
        self.d_conv: int = d_conv
        self.expand: int = expand
        self.inner_dim: int = self.dim * expand
        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))

        if is_decoder is not None:
            self.is_decoder: bool = is_decoder
        elif config is not None and not isinstance(config, int):
            self.is_decoder = bool(getattr(config, "is_decoder", False))
        else:
            self.is_decoder = False

        self.is_nested_encoder: bool = False

        self.in_proj = LoRALinear(
            self.dim,
            self.inner_dim * 2,
            bias=False,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.conv1d = nn.Conv1d(
            self.inner_dim,
            self.inner_dim,
            kernel_size=d_conv,
            groups=self.inner_dim,
            padding=d_conv - 1,
        )
        # A_log init: log([1, 2, ..., d_state]) repeated.
        # In our scan, we compute A = +exp(A_log) (positive),
        # then log_decay = -dt * A (negative), giving decay < 1.
        # This is mathematically equivalent to Mamba's
        # A = -exp(A_log), decay = exp(dt * A).
        A_init = torch.arange(1, d_state + 1, dtype=torch.float32).repeat(self.inner_dim, 1)
        self.A_log = nn.Parameter(torch.log(A_init))

        # Trainable skip connection D matching inner_dim
        self.D = nn.Parameter(torch.ones(self.inner_dim))

        # Mamba paper: dt_rank = ceil(d_model / 16), not d_inner
        self.dt_rank = math.ceil(self.dim / 16)

        # x_proj maps inner_dim → dt_rank + 2*d_state
        self.x_proj = nn.Linear(
            self.inner_dim,
            self.dt_rank + 2 * self.d_state,
            bias=False,
        )
        self.x_proj.weight.requires_grad = True

        # dt_proj maps dt_rank → inner_dim
        self.dt_proj = nn.Linear(self.dt_rank, self.inner_dim, bias=True)

        # Initialize dt_proj bias (Mamba paper Eq. 12)
        dt_min, dt_max = 0.001, 0.1
        dt_init = torch.exp(
            torch.rand(self.inner_dim) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)
        )
        dt_init_inv = dt_init + torch.log(-torch.expm1(-dt_init))
        with torch.no_grad():
            self.dt_proj.bias.copy_(dt_init_inv)
            # Weight init: uniform(-1/sqrt(dt_rank), 1/sqrt(dt_rank))
            bound = 1.0 / math.sqrt(self.dt_rank)
            nn.init.uniform_(self.dt_proj.weight, -bound, bound)
        self.dt_proj.weight.requires_grad = True
        self.dt_proj.bias.requires_grad = True

        self.out_proj = LoRALinear(
            self.inner_dim,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        if src_attn is not None:
            self.copy_weights(src_attn)

    @property
    def query(self) -> LoRALinear:
        return self.in_proj

    @property
    def key(self) -> LoRALinear:
        return self.in_proj

    @property
    def value(self) -> LoRALinear:
        return self.in_proj

    @property
    def q_lin(self) -> LoRALinear:
        return self.in_proj

    @property
    def k_lin(self) -> LoRALinear:
        return self.in_proj

    @property
    def v_lin(self) -> LoRALinear:
        return self.in_proj

    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, _b_q = None, None
            w_k, _b_k = None, None
            w_v, _b_v = None, None
            w_o, b_o = None, None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight.detach()
                    b_o = (
                        old_attn.output.dense.bias.detach()
                        if old_attn.output.dense.bias is not None
                        else None
                    )
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                (src.c_attn.bias.detach() if src.c_attn.bias is not None else None)
                if c_attn_w.dim() == 2 and (
                    c_attn_w.shape[0] == self.dim and c_attn_w.shape[1] != self.dim
                ):
                    c_attn_w = c_attn_w.t()

                split_q = min(self.dim, c_attn_w.shape[0] // 3)
                w_q = c_attn_w[:split_q, :]
                w_k = c_attn_w[split_q : 2 * split_q, :]
                w_v = c_attn_w[2 * split_q : 3 * split_q, :]

                out_w = src.c_proj.weight.detach()
                if out_w.dim() == 2 and (out_w.shape[0] != self.dim and out_w.shape[1] == self.dim):
                    out_w = out_w.t()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                src.value.bias.detach() if src.value.bias is not None else None

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                src.v_proj.bias.detach() if src.v_proj.bias is not None else None

                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # Pattern 4: DistilBERT
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                src.v_lin.bias.detach() if src.v_lin.bias is not None else None

                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # Fallback: UniversalAdapter extraction
            if w_q is None or w_v is None or (w_o is None and not self.is_nested_encoder):
                for proj in ["q", "k", "v", "o"]:
                    try:
                        w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
                        if w is not None:
                            if proj == "q" and w_q is None:
                                w_q, _b_q = w, b
                            elif proj == "k" and w_k is None:
                                w_k, _b_k = w, b
                            elif proj == "v" and w_v is None:
                                w_v, _b_v = w, b
                            elif proj == "o" and w_o is None:
                                w_o, b_o = w, b
                    except Exception as e:
                        logger.debug("UniversalAdapter extraction for %s failed: %s", proj, e)

            missing = []
            if w_q is None:
                missing.append("Q")
            if w_v is None:
                missing.append("V")
            if w_o is None and not self.is_nested_encoder:
                missing.append("O")
            if missing:
                raise RuntimeError(
                    f"SelectiveMamba: cannot extract {missing} from "
                    f"{old_attn.__class__.__name__}. Aborting to avoid "
                    f"silent model corruption."
                )

            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()
                return w

            def _adapt_to_inner(w: torch.Tensor, name: str) -> torch.Tensor:
                """
                Adapt weight from source head count to inner_dim.
                If source is smaller: tile and truncate.
                If source is larger: slice to fit.
                """
                src_out, src_in = w.shape
                if src_out == self.inner_dim:
                    return w

                if src_out > self.inner_dim:
                    # Source has more heads — take the first inner_dim
                    logger.debug(
                        "SelectiveMamba: slicing %s from %d to %d",
                        name,
                        src_out,
                        self.inner_dim,
                    )
                    return w[: self.inner_dim, :].contiguous()

                # Source is smaller — repeat to fill
                rep = math.ceil(self.inner_dim / src_out)
                logger.debug(
                    "SelectiveMamba: repeating %s from %d to %d (rep=%d)",
                    name,
                    src_out,
                    self.inner_dim,
                    rep,
                )
                return w.repeat(rep, 1)[: self.inner_dim, :].contiguous()

            w_v_or = _orient(w_v, self.dim)
            w_q_or = _orient(w_q, self.dim)
            w_v_or = _adapt_to_inner(w_v_or, "W_V")
            w_q_or = _adapt_to_inner(w_q_or, "W_Q")

            v_chunk = w_v_or[:, : self.dim]
            q_chunk = w_q_or[:, : self.dim]

            self.in_proj.weight.data[: self.inner_dim, : self.dim].copy_(v_chunk)
            self.in_proj.weight.data[self.inner_dim : 2 * self.inner_dim, : self.dim].copy_(q_chunk)

            if w_o is not None:
                w_o_or = _orient(w_o, self.inner_dim)
                min_o_out = min(self.dim, w_o_or.shape[0])
                min_o_in = min(self.dim, w_o_or.shape[1])
                self.out_proj.weight.data[:min_o_out, :min_o_in].copy_(
                    w_o_or[:min_o_out, :min_o_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))]
                    )

            # Frobenius norm checks
            for name, w in [
                ("in_proj", self.in_proj.weight),
                ("out_proj", self.out_proj.weight),
            ]:
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    logger.warning(f"SelectiveMamba: {name} has near-zero Frobenius norm (< 1e-4)")

            self.in_proj.freeze_base_weights()
            self.out_proj.freeze_base_weights()
            self.x_proj.weight.requires_grad = True
            self.dt_proj.weight.requires_grad = True
            self.dt_proj.bias.requires_grad = True

            # Preserve source LoRA adapters
            source_submodules = {
                "in": ["q_proj", "q_lin", "query", "W_q"],
                "out": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
            }
            for name in source_submodules["in"]:
                if hasattr(old_attn, name):
                    sub = getattr(old_attn, name)
                    if (
                        hasattr(sub, "lora_A")
                        and hasattr(sub, "lora_B")
                        and sub.lora_A is not None
                        and sub.lora_B is not None
                    ):
                        src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                        src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                        if hasattr(src_lA, "data") and self.in_proj.lora_A is not None:
                            r_min = min(self.in_proj.lora_A.shape[0], src_lA.shape[0])
                            in_min = min(self.in_proj.lora_A.shape[1], src_lA.shape[1])
                            self.in_proj.lora_A.data[:r_min, :in_min].copy_(
                                src_lA.data[:r_min, :in_min]
                            )
                            self.in_proj.lora_A.requires_grad = True
                        if hasattr(src_lB, "data") and self.in_proj.lora_B is not None:
                            out_min = min(self.in_proj.lora_B.shape[0], src_lB.shape[0])
                            r_min = min(self.in_proj.lora_B.shape[1], src_lB.shape[1])
                            self.in_proj.lora_B.data[:out_min, :r_min].copy_(
                                src_lB.data[:out_min, :r_min]
                            )
                            self.in_proj.lora_B.requires_grad = True
                        if hasattr(sub, "lora_scaling"):
                            self.in_proj.scaling = sub.lora_scaling
                        break

            for name in source_submodules["out"]:
                if hasattr(old_attn, name):
                    sub = getattr(old_attn, name)
                    if (
                        hasattr(sub, "lora_A")
                        and hasattr(sub, "lora_B")
                        and sub.lora_A is not None
                        and sub.lora_B is not None
                    ):
                        src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                        src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                        if hasattr(src_lA, "data") and self.out_proj.lora_A is not None:
                            r_min = min(self.out_proj.lora_A.shape[0], src_lA.shape[0])
                            in_min = min(self.out_proj.lora_A.shape[1], src_lA.shape[1])
                            self.out_proj.lora_A.data[:r_min, :in_min].copy_(
                                src_lA.data[:r_min, :in_min]
                            )
                            self.out_proj.lora_A.requires_grad = True
                        if hasattr(src_lB, "data") and self.out_proj.lora_B is not None:
                            out_min = min(self.out_proj.lora_B.shape[0], src_lB.shape[0])
                            r_min = min(self.out_proj.lora_B.shape[1], src_lB.shape[1])
                            self.out_proj.lora_B.data[:out_min, :r_min].copy_(
                                src_lB.data[:out_min, :r_min]
                            )
                            self.out_proj.lora_B.requires_grad = True
                        if hasattr(sub, "lora_scaling"):
                            self.out_proj.scaling = sub.lora_scaling
                        break

    def _selective_scan(
        self,
        x: torch.Tensor,
        dt: torch.Tensor,
        B_t: torch.Tensor,
        C_t: torch.Tensor,
        mask: torch.Tensor | None = None,
        reverse: bool = False,
    ) -> torch.Tensor:
        """Bounded chunked associative SSM scan in FP32 with O(chunk=64) memory."""
        if reverse:
            x = torch.flip(x, dims=[1])
            dt = torch.flip(dt, dims=[1])
            B_t = torch.flip(B_t, dims=[1])
            C_t = torch.flip(C_t, dims=[1])
            if mask is not None:
                mask = torch.flip(mask, dims=[1])

        batch, seqlen, _ = x.shape
        A = torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0))  # [inner_dim, d_state]

        # Path 1: Hardware-fused Triton selective scan on CUDA
        if HAS_TRITON_UPI and triton_selective_scan is not None and x.is_cuda and mask is None:
            try:
                return triton_selective_scan(
                    x.contiguous(),
                    dt.contiguous(),
                    -A.contiguous(),
                    B_t.contiguous(),
                    C_t.contiguous(),
                    None,
                    delta_softplus=False,
                ).to(x.dtype)
            except Exception as e:
                logger.debug("Triton selective scan failed, falling back to chunked PyTorch: %s", e)

        # In case of padding mask, zero out invalid token inputs
        if mask is not None:
            m = mask.unsqueeze(-1).float()
            x = x * m
            dt = dt * m

        chunk_size = 64
        num_chunks = (seqlen + chunk_size - 1) // chunk_size
        h_running = x.new_zeros(batch, self.inner_dim, self.d_state, dtype=torch.float32)
        y_chunks = []
        A_view = A.view(1, 1, self.inner_dim, self.d_state)

        for c in range(num_chunks):
            s_idx = c * chunk_size
            e_idx = min(seqlen, s_idx + chunk_size)

            dt_c = dt[:, s_idx:e_idx]
            x_c = x[:, s_idx:e_idx]
            B_c = B_t[:, s_idx:e_idx]
            C_c = C_t[:, s_idx:e_idx]

            # Compute u and log_a strictly for current 64-token chunk to avoid buffer overflow
            log_a_c = -(dt_c.unsqueeze(-1) * A_view).float()
            u_c = (dt_c.unsqueeze(-1) * B_c.unsqueeze(2) * x_c.unsqueeze(-1)).float()

            lambda_c = torch.clamp(torch.cumsum(log_a_c, dim=1), min=-25.0, max=0.0)

            decay_c = torch.exp(lambda_c)
            decay_inv_c = torch.exp(-lambda_c).clamp(max=1e10)

            h_intra = decay_c * torch.cumsum(decay_inv_c * u_c, dim=1)
            h_inter = decay_c * h_running.unsqueeze(1)

            h_chunk = torch.clamp(h_intra + h_inter, min=-1e4, max=1e4)
            y_c = (C_c.unsqueeze(2).float() * h_chunk).sum(dim=-1).to(x.dtype)
            y_chunks.append(y_c)
            h_running = h_chunk[:, -1]

        y = torch.cat(y_chunks, dim=1)
        y = (y + self.D.to(x.dtype) * x).to(x.dtype)

        if reverse:
            y = torch.flip(y, dims=[1])

        return y

    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]:
        if encoder_hidden_states is not None:
            raise NotImplementedError(
                "SelectiveMamba does not support cross-attention. "
                "This module replaces SELF-attention only. "
                "For cross-attention, use a different module."
            )

        B, L, D = hidden_states.shape
        mask = attention_mask if attention_mask is not None else kwargs.get("mask")

        if head_mask is not None:
            logger.warning(
                "SelectiveMamba does not use head_mask (SSM has no heads). Ignoring head_mask."
            )

        mask_bool = None
        if mask is not None:
            if mask.dim() == 2:
                # Distinguish between int/bool masks and float additive masks
                if mask.dtype in (
                    torch.bool,
                    torch.int8,
                    torch.int16,
                    torch.int32,
                    torch.int64,
                ):
                    mask_bool = mask.bool().to(hidden_states.device)
                else:
                    # Float mask: -inf = invalid, 0 = valid
                    mask_bool = (mask > -1e4).to(hidden_states.device)
            elif mask.dim() == 4:
                mask_bool = (mask[:, 0, 0, :] > -1e4).to(hidden_states.device)
            else:
                try:
                    mask_bool = get_token_valid_mask(
                        mask,
                        B,
                        L,
                        hidden_states.device,
                    )
                except Exception as e:
                    logger.debug("get_token_valid_mask failed: %s", e)
                    mask_bool = None

        target_dtype = hidden_states.dtype
        xz = self.in_proj(hidden_states)
        x_in, z = xz.chunk(2, dim=-1)

        conv_w = self.conv1d.weight.to(x_in.dtype)
        conv_b = self.conv1d.bias.to(x_in.dtype) if self.conv1d.bias is not None else None
        x_conv = (
            F.conv1d(
                x_in.transpose(1, 2),
                conv_w,
                conv_b,
                stride=self.conv1d.stride,
                padding=self.conv1d.padding,
                dilation=self.conv1d.dilation,
                groups=self.conv1d.groups,
            )[:, :, :L]
            .transpose(1, 2)
            .contiguous()
        )
        x_ssm = F.silu(x_conv)

        # Selective projections: dt_raw [dt_rank], B_t [d_state], C_t [d_state]
        x_proj_out = self.x_proj(x_ssm)
        dt_raw, B_t, C_t = torch.split(
            x_proj_out,
            [self.dt_rank, self.d_state, self.d_state],
            dim=-1,
        )
        dt = F.softplus(self.dt_proj(dt_raw))
        dt = torch.clamp(dt, min=1e-3, max=1.0)

        mt = str(getattr(getattr(self, "config", None), "model_type", "") or "").lower()
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_decoder = (
            self.is_decoder
            or getattr(getattr(self, "config", None), "is_decoder", False)
            or any(mt.startswith(p) for p in decoder_prefixes)
        )

        y_fwd = self._selective_scan(x_ssm, dt, B_t, C_t, mask=mask_bool, reverse=False)

        if is_decoder:
            y = y_fwd
        else:
            y_bwd = self._selective_scan(x_ssm, dt, B_t, C_t, mask=mask_bool, reverse=True)
            y = 0.5 * (y_fwd + y_bwd)

        gated = (y.to(target_dtype) * F.silu(z.to(target_dtype))).to(target_dtype)

        # Inner SSM output is inner_dim; out_proj brings it back to dim.
        # Always apply out_proj.
        out = self.out_proj(gated).to(target_dtype)

        if mask_bool is not None:
            out = out * mask_bool.unsqueeze(-1).to(out.dtype)

        return (out, None)

    def step_recurrent(
        self,
        x_t: torch.Tensor,
        state: dict[str, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """
        Single-step streaming inference.

        x_t: (B, 1, D)
        state: dict with keys 'h', 'conv_buf' (optional)
        Returns: (y_t, new_state)
        """
        if x_t.dim() == 2:
            x_t = x_t.unsqueeze(1)
        B, L, D = x_t.shape
        assert L == 1, "step_recurrent expects L=1"

        # Initialize state
        if state is None:
            h = torch.zeros(
                B,
                self.inner_dim,
                self.d_state,
                device=x_t.device,
                dtype=torch.float32,
            )
            conv_buf = torch.zeros(
                B,
                self.d_conv - 1,
                self.inner_dim,
                device=x_t.device,
                dtype=x_t.dtype,
            )
        else:
            h = state["h"]
            conv_buf = state.get("conv_buf")

        # in_proj
        xz = self.in_proj(x_t)  # (B, 1, 2*inner_dim)
        x_in, z = xz.chunk(2, dim=-1)

        # Conv1d with buffer
        x_conv_input = torch.cat([conv_buf, x_in], dim=1)  # (B, d_conv, inner_dim)
        x_conv = F.conv1d(
            x_conv_input.transpose(1, 2),
            self.conv1d.weight,
            self.conv1d.bias,
            groups=self.inner_dim,
        ).transpose(1, 2)  # (B, 1, inner_dim)
        new_conv_buf = x_conv_input[:, -(self.d_conv - 1) :]

        x_ssm = F.silu(x_conv)

        # Selective params
        x_proj_out = self.x_proj(x_ssm)
        dt_raw, B_t, C_t = torch.split(
            x_proj_out,
            [self.dt_rank, self.d_state, self.d_state],
            dim=-1,
        )
        dt = F.softplus(self.dt_proj(dt_raw)).clamp(1e-3, 1.0)

        # Single step scan
        A = torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0))
        dA = torch.exp(-dt.unsqueeze(-1) * A.view(1, 1, self.inner_dim, self.d_state))
        dBx = dt.unsqueeze(-1) * B_t.unsqueeze(2) * x_ssm.unsqueeze(-1)

        h = dA.squeeze(1) * h + dBx.squeeze(1)  # (B, inner_dim, d_state)
        y = (C_t.squeeze(1).unsqueeze(1) * h).sum(dim=-1)  # (B, inner_dim)
        y = y + self.D * x_ssm.squeeze(1)

        # Gating
        gated = y * F.silu(z.squeeze(1))  # (B, inner_dim)

        # Output projection
        out = self.out_proj(gated.unsqueeze(1))  # (B, 1, D)

        new_state = {"h": h, "conv_buf": new_conv_buf}
        return out, new_state

Selective Mamba SSM block.

Implements a Mamba-style selective state space model with input-dependent B, C, and dt parameters: h_t = Ā_t · h_{t-1} + B̄_t · x_t y_t = C_t · h_t + D · x_t where Ā_t = exp(dt_t · A).

The "selectivity" comes from x_proj producing B_t and C_t from the input.

NOTE: The previous name "UPI-Stuffed" has been removed because none of the claimed features (Uncertainty Penalty, Interpolation, Stuffing) were actually implemented.

Registered as "upi_mamba" (legacy alias) and "selective_mamba" (canonical).

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.

Instance variables

prop k_lin : LoRALinear

Expand source code

@property
def k_lin(self) -> LoRALinear:
    return self.in_proj

prop key : LoRALinear

Expand source code

@property
def key(self) -> LoRALinear:
    return self.in_proj

prop q_lin : LoRALinear

Expand source code

@property
def q_lin(self) -> LoRALinear:
    return self.in_proj

prop query : LoRALinear

Expand source code

@property
def query(self) -> LoRALinear:
    return self.in_proj

prop v_lin : LoRALinear

Expand source code

@property
def v_lin(self) -> LoRALinear:
    return self.in_proj

prop value : LoRALinear

Expand source code

@property
def value(self) -> LoRALinear:
    return self.in_proj

Methods

def copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, _b_q = None, None
        w_k, _b_k = None, None
        w_v, _b_v = None, None
        w_o, b_o = None, None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight.detach()
                b_o = (
                    old_attn.output.dense.bias.detach()
                    if old_attn.output.dense.bias is not None
                    else None
                )
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            (src.c_attn.bias.detach() if src.c_attn.bias is not None else None)
            if c_attn_w.dim() == 2 and (
                c_attn_w.shape[0] == self.dim and c_attn_w.shape[1] != self.dim
            ):
                c_attn_w = c_attn_w.t()

            split_q = min(self.dim, c_attn_w.shape[0] // 3)
            w_q = c_attn_w[:split_q, :]
            w_k = c_attn_w[split_q : 2 * split_q, :]
            w_v = c_attn_w[2 * split_q : 3 * split_q, :]

            out_w = src.c_proj.weight.detach()
            if out_w.dim() == 2 and (out_w.shape[0] != self.dim and out_w.shape[1] == self.dim):
                out_w = out_w.t()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            src.value.bias.detach() if src.value.bias is not None else None

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            src.v_proj.bias.detach() if src.v_proj.bias is not None else None

            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # Pattern 4: DistilBERT
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            src.v_lin.bias.detach() if src.v_lin.bias is not None else None

            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # Fallback: UniversalAdapter extraction
        if w_q is None or w_v is None or (w_o is None and not self.is_nested_encoder):
            for proj in ["q", "k", "v", "o"]:
                try:
                    w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
                    if w is not None:
                        if proj == "q" and w_q is None:
                            w_q, _b_q = w, b
                        elif proj == "k" and w_k is None:
                            w_k, _b_k = w, b
                        elif proj == "v" and w_v is None:
                            w_v, _b_v = w, b
                        elif proj == "o" and w_o is None:
                            w_o, b_o = w, b
                except Exception as e:
                    logger.debug("UniversalAdapter extraction for %s failed: %s", proj, e)

        missing = []
        if w_q is None:
            missing.append("Q")
        if w_v is None:
            missing.append("V")
        if w_o is None and not self.is_nested_encoder:
            missing.append("O")
        if missing:
            raise RuntimeError(
                f"SelectiveMamba: cannot extract {missing} from "
                f"{old_attn.__class__.__name__}. Aborting to avoid "
                f"silent model corruption."
            )

        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()
            return w

        def _adapt_to_inner(w: torch.Tensor, name: str) -> torch.Tensor:
            """
            Adapt weight from source head count to inner_dim.
            If source is smaller: tile and truncate.
            If source is larger: slice to fit.
            """
            src_out, src_in = w.shape
            if src_out == self.inner_dim:
                return w

            if src_out > self.inner_dim:
                # Source has more heads — take the first inner_dim
                logger.debug(
                    "SelectiveMamba: slicing %s from %d to %d",
                    name,
                    src_out,
                    self.inner_dim,
                )
                return w[: self.inner_dim, :].contiguous()

            # Source is smaller — repeat to fill
            rep = math.ceil(self.inner_dim / src_out)
            logger.debug(
                "SelectiveMamba: repeating %s from %d to %d (rep=%d)",
                name,
                src_out,
                self.inner_dim,
                rep,
            )
            return w.repeat(rep, 1)[: self.inner_dim, :].contiguous()

        w_v_or = _orient(w_v, self.dim)
        w_q_or = _orient(w_q, self.dim)
        w_v_or = _adapt_to_inner(w_v_or, "W_V")
        w_q_or = _adapt_to_inner(w_q_or, "W_Q")

        v_chunk = w_v_or[:, : self.dim]
        q_chunk = w_q_or[:, : self.dim]

        self.in_proj.weight.data[: self.inner_dim, : self.dim].copy_(v_chunk)
        self.in_proj.weight.data[self.inner_dim : 2 * self.inner_dim, : self.dim].copy_(q_chunk)

        if w_o is not None:
            w_o_or = _orient(w_o, self.inner_dim)
            min_o_out = min(self.dim, w_o_or.shape[0])
            min_o_in = min(self.dim, w_o_or.shape[1])
            self.out_proj.weight.data[:min_o_out, :min_o_in].copy_(
                w_o_or[:min_o_out, :min_o_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))]
                )

        # Frobenius norm checks
        for name, w in [
            ("in_proj", self.in_proj.weight),
            ("out_proj", self.out_proj.weight),
        ]:
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                logger.warning(f"SelectiveMamba: {name} has near-zero Frobenius norm (< 1e-4)")

        self.in_proj.freeze_base_weights()
        self.out_proj.freeze_base_weights()
        self.x_proj.weight.requires_grad = True
        self.dt_proj.weight.requires_grad = True
        self.dt_proj.bias.requires_grad = True

        # Preserve source LoRA adapters
        source_submodules = {
            "in": ["q_proj", "q_lin", "query", "W_q"],
            "out": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
        }
        for name in source_submodules["in"]:
            if hasattr(old_attn, name):
                sub = getattr(old_attn, name)
                if (
                    hasattr(sub, "lora_A")
                    and hasattr(sub, "lora_B")
                    and sub.lora_A is not None
                    and sub.lora_B is not None
                ):
                    src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                    src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                    if hasattr(src_lA, "data") and self.in_proj.lora_A is not None:
                        r_min = min(self.in_proj.lora_A.shape[0], src_lA.shape[0])
                        in_min = min(self.in_proj.lora_A.shape[1], src_lA.shape[1])
                        self.in_proj.lora_A.data[:r_min, :in_min].copy_(
                            src_lA.data[:r_min, :in_min]
                        )
                        self.in_proj.lora_A.requires_grad = True
                    if hasattr(src_lB, "data") and self.in_proj.lora_B is not None:
                        out_min = min(self.in_proj.lora_B.shape[0], src_lB.shape[0])
                        r_min = min(self.in_proj.lora_B.shape[1], src_lB.shape[1])
                        self.in_proj.lora_B.data[:out_min, :r_min].copy_(
                            src_lB.data[:out_min, :r_min]
                        )
                        self.in_proj.lora_B.requires_grad = True
                    if hasattr(sub, "lora_scaling"):
                        self.in_proj.scaling = sub.lora_scaling
                    break

        for name in source_submodules["out"]:
            if hasattr(old_attn, name):
                sub = getattr(old_attn, name)
                if (
                    hasattr(sub, "lora_A")
                    and hasattr(sub, "lora_B")
                    and sub.lora_A is not None
                    and sub.lora_B is not None
                ):
                    src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                    src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                    if hasattr(src_lA, "data") and self.out_proj.lora_A is not None:
                        r_min = min(self.out_proj.lora_A.shape[0], src_lA.shape[0])
                        in_min = min(self.out_proj.lora_A.shape[1], src_lA.shape[1])
                        self.out_proj.lora_A.data[:r_min, :in_min].copy_(
                            src_lA.data[:r_min, :in_min]
                        )
                        self.out_proj.lora_A.requires_grad = True
                    if hasattr(src_lB, "data") and self.out_proj.lora_B is not None:
                        out_min = min(self.out_proj.lora_B.shape[0], src_lB.shape[0])
                        r_min = min(self.out_proj.lora_B.shape[1], src_lB.shape[1])
                        self.out_proj.lora_B.data[:out_min, :r_min].copy_(
                            src_lB.data[:out_min, :r_min]
                        )
                        self.out_proj.lora_B.requires_grad = True
                    if hasattr(sub, "lora_scaling"):
                        self.out_proj.scaling = sub.lora_scaling
                    break

Robust weight copying from DistilBERT, RoBERTa, GPT-2 Conv1D, or Qwen/LLaMA.

def step_recurrent(self, x_t: torch.Tensor, state: dict[str, torch.Tensor] | None = None) ‑> tuple[torch.Tensor, dict[str, torch.Tensor]]

Expand source code

def step_recurrent(
    self,
    x_t: torch.Tensor,
    state: dict[str, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """
    Single-step streaming inference.

    x_t: (B, 1, D)
    state: dict with keys 'h', 'conv_buf' (optional)
    Returns: (y_t, new_state)
    """
    if x_t.dim() == 2:
        x_t = x_t.unsqueeze(1)
    B, L, D = x_t.shape
    assert L == 1, "step_recurrent expects L=1"

    # Initialize state
    if state is None:
        h = torch.zeros(
            B,
            self.inner_dim,
            self.d_state,
            device=x_t.device,
            dtype=torch.float32,
        )
        conv_buf = torch.zeros(
            B,
            self.d_conv - 1,
            self.inner_dim,
            device=x_t.device,
            dtype=x_t.dtype,
        )
    else:
        h = state["h"]
        conv_buf = state.get("conv_buf")

    # in_proj
    xz = self.in_proj(x_t)  # (B, 1, 2*inner_dim)
    x_in, z = xz.chunk(2, dim=-1)

    # Conv1d with buffer
    x_conv_input = torch.cat([conv_buf, x_in], dim=1)  # (B, d_conv, inner_dim)
    x_conv = F.conv1d(
        x_conv_input.transpose(1, 2),
        self.conv1d.weight,
        self.conv1d.bias,
        groups=self.inner_dim,
    ).transpose(1, 2)  # (B, 1, inner_dim)
    new_conv_buf = x_conv_input[:, -(self.d_conv - 1) :]

    x_ssm = F.silu(x_conv)

    # Selective params
    x_proj_out = self.x_proj(x_ssm)
    dt_raw, B_t, C_t = torch.split(
        x_proj_out,
        [self.dt_rank, self.d_state, self.d_state],
        dim=-1,
    )
    dt = F.softplus(self.dt_proj(dt_raw)).clamp(1e-3, 1.0)

    # Single step scan
    A = torch.exp(torch.clamp(self.A_log.float(), min=-20.0, max=10.0))
    dA = torch.exp(-dt.unsqueeze(-1) * A.view(1, 1, self.inner_dim, self.d_state))
    dBx = dt.unsqueeze(-1) * B_t.unsqueeze(2) * x_ssm.unsqueeze(-1)

    h = dA.squeeze(1) * h + dBx.squeeze(1)  # (B, inner_dim, d_state)
    y = (C_t.squeeze(1).unsqueeze(1) * h).sum(dim=-1)  # (B, inner_dim)
    y = y + self.D * x_ssm.squeeze(1)

    # Gating
    gated = y * F.silu(z.squeeze(1))  # (B, inner_dim)

    # Output projection
    out = self.out_proj(gated.unsqueeze(1))  # (B, 1, D)

    new_state = {"h": h, "conv_buf": new_conv_buf}
    return out, new_state

Single-step streaming inference.

x_t: (B, 1, D) state: dict with keys 'h', 'conv_buf' (optional) Returns: (y_t, new_state)

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class SelectiveSSM (config: Any | None = None,
rank: int | None = None,
d_model: int | None = None,
d_state: int = 16,
d_conv: int = 4,
expand: int = 2,
dt_rank: int | None = None,
conv_bias: bool = True,
bias: bool = False,
bidirectional: bool = False,
lora_rank: int = 0,
weights_init: str = 'mamba',
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("selective_ssm")
@ATTENTION_REGISTRY.register("mamba_ssm")
class SelectiveSSM(BaseAttention):
    r"""
    Faithful Mamba-style selective SSM block.
    """

    supports_causal: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "state_update"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | None = None,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 2,
        dt_rank: int | None = None,
        conv_bias: bool = True,
        bias: bool = False,
        bidirectional: bool = False,
        lora_rank: int = 0,
        weights_init: str = "mamba",
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        if d_model is not None:
            self.d_model = int(d_model)
        else:
            self.d_model = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))

        self.d_state = int(_cfg_get(config, "d_state", kwargs.get("d_state", d_state)))
        self.d_conv = int(_cfg_get(config, "d_conv", kwargs.get("d_conv", d_conv)))
        self.expand = int(_cfg_get(config, "expand", kwargs.get("expand", expand)))
        self.d_inner = int(self.expand * self.d_model)
        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank is None else int(dt_rank)
        self.bidirectional = bool(
            _cfg_get(config, "bidirectional", kwargs.get("bidirectional", bidirectional))
        )
        self.weights_init = weights_init

        # Input projection: x -> (x_ssm, z)
        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias)

        # Depthwise 1D conv
        self.conv1d = nn.Conv1d(
            self.d_inner,
            self.d_inner,
            kernel_size=self.d_conv,
            groups=self.d_inner,
            padding=self.d_conv - 1,
            bias=conv_bias,
        )

        # Delta, B, C projections from x_ssm
        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + 2 * self.d_state, bias=False)

        # Delta projection from dt_rank -> d_inner
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)

        # A parameterization: A = -exp(A_log)
        A = torch.arange(1, self.d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))

        # D skip connection
        self.D = nn.Parameter(torch.ones(self.d_inner))

        # Output projection
        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias)

        # Bidirectional components if requested
        if self.bidirectional:
            self.conv1d_bwd = nn.Conv1d(
                self.d_inner,
                self.d_inner,
                kernel_size=self.d_conv,
                groups=self.d_inner,
                padding=self.d_conv - 1,
                bias=conv_bias,
            )
            self.x_proj_bwd = nn.Linear(self.d_inner, self.dt_rank + 2 * self.d_state, bias=False)
            self.dt_proj_bwd = nn.Linear(self.dt_rank, self.d_inner, bias=True)
            A_bwd = torch.arange(1, self.d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)
            self.A_log_bwd = nn.Parameter(torch.log(A_bwd))
            self.D_bwd = nn.Parameter(torch.ones(self.d_inner))
            self.out_proj_bwd = nn.Linear(self.d_inner, self.d_model, bias=bias)
        self.triton_kwargs = {}

        self._reset_parameters()

    def _reset_parameters(self) -> None:
        if self.weights_init == "mamba":
            # Delta bias init: softplus^-1(Uniform(0.001, 0.1))
            dt_min, dt_max = 0.001, 0.1
            dt_init = torch.exp(
                torch.rand(self.d_inner) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)
            )
            dt_init_inv = dt_init + torch.log(-torch.expm1(-dt_init))
            with torch.no_grad():
                self.dt_proj.bias.copy_(dt_init_inv)
                self.dt_proj.weight.copy_(
                    torch.ones(self.d_inner, self.dt_rank) * (0.5 / self.dt_rank) ** 0.5
                )
                if self.bidirectional:
                    self.dt_proj_bwd.bias.copy_(dt_init_inv)
                    self.dt_proj_bwd.weight.copy_(
                        torch.ones(self.d_inner, self.dt_rank) * (0.5 / self.dt_rank) ** 0.5
                    )
        elif self.weights_init == "xavier":
            nn.init.xavier_uniform_(self.in_proj.weight)
            nn.init.xavier_uniform_(self.out_proj.weight)
            nn.init.xavier_uniform_(self.x_proj.weight)

    def _forward_core(
        self,
        x_ssm: torch.Tensor,
        z: torch.Tensor,
        conv1d_layer: nn.Module,
        x_proj_layer: nn.Module,
        dt_proj_layer: nn.Module,
        A_log_param: nn.Parameter,
        D_param: nn.Parameter,
    ) -> torch.Tensor:
        B, L, _ = x_ssm.shape

        # Conv1d
        x_conv = conv1d_layer(x_ssm.transpose(1, 2))[..., :L].transpose(1, 2)
        x_conv = F.silu(x_conv)

        # Selective projections
        x_dbl = x_proj_layer(x_conv)  # (B, L, dt_rank + 2 * d_state)
        dt, B_ssm, C_ssm = x_dbl.split([self.dt_rank, self.d_state, self.d_state], dim=-1)
        dt = dt_proj_layer(dt)  # (B, L, d_inner)
        dt = F.softplus(dt)

        A = -torch.exp(A_log_param.float())  # (d_inner, d_state)

        y = selective_scan(
            u=x_conv.transpose(1, 2),
            delta=dt.transpose(1, 2),
            A=A,
            B=B_ssm.transpose(1, 2),
            C=C_ssm.transpose(1, 2),
            D=D_param,
        )  # (B, d_inner, L)
        y = y.transpose(1, 2)  # (B, L, d_inner)

        # Gating
        target_dtype = x_ssm.dtype
        y = (y.to(target_dtype) * F.silu(z.to(target_dtype))).to(target_dtype)
        return y

    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]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("SelectiveSSM does not support cross-attention.")

        target_dtype = hidden_states.dtype
        B, L, D = hidden_states.shape

        # Project directly into SSM and gating paths without duplicating LayerNorm
        x_proj = self.in_proj(hidden_states).to(target_dtype)
        x_ssm, z = x_proj.chunk(2, dim=-1)

        y_fwd = self._forward_core(
            x_ssm=x_ssm,
            z=z,
            conv1d_layer=self.conv1d,
            x_proj_layer=self.x_proj,
            dt_proj_layer=self.dt_proj,
            A_log_param=self.A_log,
            D_param=self.D,
        )
        out = self.out_proj(y_fwd.to(target_dtype)).to(target_dtype)

        if self.bidirectional:
            x_ssm_bwd = torch.flip(x_ssm, [1])
            z_bwd = torch.flip(z, [1])
            y_bwd = self._forward_core(
                x_ssm=x_ssm_bwd,
                z=z_bwd,
                conv1d_layer=self.conv1d_bwd,
                x_proj_layer=self.x_proj_bwd,
                dt_proj_layer=self.dt_proj_bwd,
                A_log_param=self.A_log_bwd,
                D_param=self.D_bwd,
            )
            y_bwd = torch.flip(y_bwd, [1])
            out = (out + self.out_proj_bwd(y_bwd.to(target_dtype))).to(target_dtype)

        return (out.to(target_dtype), None)

Faithful Mamba-style selective SSM block.

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 triton_kernel : str

The type of the None singleton.

Methods

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]

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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    if encoder_hidden_states is not None:
        raise NotImplementedError("SelectiveSSM does not support cross-attention.")

    target_dtype = hidden_states.dtype
    B, L, D = hidden_states.shape

    # Project directly into SSM and gating paths without duplicating LayerNorm
    x_proj = self.in_proj(hidden_states).to(target_dtype)
    x_ssm, z = x_proj.chunk(2, dim=-1)

    y_fwd = self._forward_core(
        x_ssm=x_ssm,
        z=z,
        conv1d_layer=self.conv1d,
        x_proj_layer=self.x_proj,
        dt_proj_layer=self.dt_proj,
        A_log_param=self.A_log,
        D_param=self.D,
    )
    out = self.out_proj(y_fwd.to(target_dtype)).to(target_dtype)

    if self.bidirectional:
        x_ssm_bwd = torch.flip(x_ssm, [1])
        z_bwd = torch.flip(z, [1])
        y_bwd = self._forward_core(
            x_ssm=x_ssm_bwd,
            z=z_bwd,
            conv1d_layer=self.conv1d_bwd,
            x_proj_layer=self.x_proj_bwd,
            dt_proj_layer=self.dt_proj_bwd,
            A_log_param=self.A_log_bwd,
            D_param=self.D_bwd,
        )
        y_bwd = torch.flip(y_bwd, [1])
        out = (out + self.out_proj_bwd(y_bwd.to(target_dtype))).to(target_dtype)

    return (out.to(target_dtype), None)

class SimpleSSM (config: Any | None = None,
rank: int | None = None,
d_model: int | None = None,
d_state: int = 16,
d_conv: int = 4,
expand: int = 2,
dt_rank: int | None = None,
conv_bias: bool = True,
bias: bool = False,
bidirectional: bool = False,
lora_rank: int = 0,
weights_init: str = 'mamba',
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("selective_ssm")
@ATTENTION_REGISTRY.register("mamba_ssm")
class SelectiveSSM(BaseAttention):
    r"""
    Faithful Mamba-style selective SSM block.
    """

    supports_causal: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "state_update"

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        d_model: int | None = None,
        d_state: int = 16,
        d_conv: int = 4,
        expand: int = 2,
        dt_rank: int | None = None,
        conv_bias: bool = True,
        bias: bool = False,
        bidirectional: bool = False,
        lora_rank: int = 0,
        weights_init: str = "mamba",
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        if d_model is not None:
            self.d_model = int(d_model)
        else:
            self.d_model = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))

        self.d_state = int(_cfg_get(config, "d_state", kwargs.get("d_state", d_state)))
        self.d_conv = int(_cfg_get(config, "d_conv", kwargs.get("d_conv", d_conv)))
        self.expand = int(_cfg_get(config, "expand", kwargs.get("expand", expand)))
        self.d_inner = int(self.expand * self.d_model)
        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank is None else int(dt_rank)
        self.bidirectional = bool(
            _cfg_get(config, "bidirectional", kwargs.get("bidirectional", bidirectional))
        )
        self.weights_init = weights_init

        # Input projection: x -> (x_ssm, z)
        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias)

        # Depthwise 1D conv
        self.conv1d = nn.Conv1d(
            self.d_inner,
            self.d_inner,
            kernel_size=self.d_conv,
            groups=self.d_inner,
            padding=self.d_conv - 1,
            bias=conv_bias,
        )

        # Delta, B, C projections from x_ssm
        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + 2 * self.d_state, bias=False)

        # Delta projection from dt_rank -> d_inner
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)

        # A parameterization: A = -exp(A_log)
        A = torch.arange(1, self.d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))

        # D skip connection
        self.D = nn.Parameter(torch.ones(self.d_inner))

        # Output projection
        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias)

        # Bidirectional components if requested
        if self.bidirectional:
            self.conv1d_bwd = nn.Conv1d(
                self.d_inner,
                self.d_inner,
                kernel_size=self.d_conv,
                groups=self.d_inner,
                padding=self.d_conv - 1,
                bias=conv_bias,
            )
            self.x_proj_bwd = nn.Linear(self.d_inner, self.dt_rank + 2 * self.d_state, bias=False)
            self.dt_proj_bwd = nn.Linear(self.dt_rank, self.d_inner, bias=True)
            A_bwd = torch.arange(1, self.d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)
            self.A_log_bwd = nn.Parameter(torch.log(A_bwd))
            self.D_bwd = nn.Parameter(torch.ones(self.d_inner))
            self.out_proj_bwd = nn.Linear(self.d_inner, self.d_model, bias=bias)
        self.triton_kwargs = {}

        self._reset_parameters()

    def _reset_parameters(self) -> None:
        if self.weights_init == "mamba":
            # Delta bias init: softplus^-1(Uniform(0.001, 0.1))
            dt_min, dt_max = 0.001, 0.1
            dt_init = torch.exp(
                torch.rand(self.d_inner) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)
            )
            dt_init_inv = dt_init + torch.log(-torch.expm1(-dt_init))
            with torch.no_grad():
                self.dt_proj.bias.copy_(dt_init_inv)
                self.dt_proj.weight.copy_(
                    torch.ones(self.d_inner, self.dt_rank) * (0.5 / self.dt_rank) ** 0.5
                )
                if self.bidirectional:
                    self.dt_proj_bwd.bias.copy_(dt_init_inv)
                    self.dt_proj_bwd.weight.copy_(
                        torch.ones(self.d_inner, self.dt_rank) * (0.5 / self.dt_rank) ** 0.5
                    )
        elif self.weights_init == "xavier":
            nn.init.xavier_uniform_(self.in_proj.weight)
            nn.init.xavier_uniform_(self.out_proj.weight)
            nn.init.xavier_uniform_(self.x_proj.weight)

    def _forward_core(
        self,
        x_ssm: torch.Tensor,
        z: torch.Tensor,
        conv1d_layer: nn.Module,
        x_proj_layer: nn.Module,
        dt_proj_layer: nn.Module,
        A_log_param: nn.Parameter,
        D_param: nn.Parameter,
    ) -> torch.Tensor:
        B, L, _ = x_ssm.shape

        # Conv1d
        x_conv = conv1d_layer(x_ssm.transpose(1, 2))[..., :L].transpose(1, 2)
        x_conv = F.silu(x_conv)

        # Selective projections
        x_dbl = x_proj_layer(x_conv)  # (B, L, dt_rank + 2 * d_state)
        dt, B_ssm, C_ssm = x_dbl.split([self.dt_rank, self.d_state, self.d_state], dim=-1)
        dt = dt_proj_layer(dt)  # (B, L, d_inner)
        dt = F.softplus(dt)

        A = -torch.exp(A_log_param.float())  # (d_inner, d_state)

        y = selective_scan(
            u=x_conv.transpose(1, 2),
            delta=dt.transpose(1, 2),
            A=A,
            B=B_ssm.transpose(1, 2),
            C=C_ssm.transpose(1, 2),
            D=D_param,
        )  # (B, d_inner, L)
        y = y.transpose(1, 2)  # (B, L, d_inner)

        # Gating
        target_dtype = x_ssm.dtype
        y = (y.to(target_dtype) * F.silu(z.to(target_dtype))).to(target_dtype)
        return y

    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]:
        if encoder_hidden_states is not None:
            raise NotImplementedError("SelectiveSSM does not support cross-attention.")

        target_dtype = hidden_states.dtype
        B, L, D = hidden_states.shape

        # Project directly into SSM and gating paths without duplicating LayerNorm
        x_proj = self.in_proj(hidden_states).to(target_dtype)
        x_ssm, z = x_proj.chunk(2, dim=-1)

        y_fwd = self._forward_core(
            x_ssm=x_ssm,
            z=z,
            conv1d_layer=self.conv1d,
            x_proj_layer=self.x_proj,
            dt_proj_layer=self.dt_proj,
            A_log_param=self.A_log,
            D_param=self.D,
        )
        out = self.out_proj(y_fwd.to(target_dtype)).to(target_dtype)

        if self.bidirectional:
            x_ssm_bwd = torch.flip(x_ssm, [1])
            z_bwd = torch.flip(z, [1])
            y_bwd = self._forward_core(
                x_ssm=x_ssm_bwd,
                z=z_bwd,
                conv1d_layer=self.conv1d_bwd,
                x_proj_layer=self.x_proj_bwd,
                dt_proj_layer=self.dt_proj_bwd,
                A_log_param=self.A_log_bwd,
                D_param=self.D_bwd,
            )
            y_bwd = torch.flip(y_bwd, [1])
            out = (out + self.out_proj_bwd(y_bwd.to(target_dtype))).to(target_dtype)

        return (out.to(target_dtype), None)

Faithful Mamba-style selective SSM block.

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 triton_kernel : str

The type of the None singleton.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class SharedTokenScorer (dim: int, hidden: int = 64)

Expand source code

class SharedTokenScorer(nn.Module):
    """
    Optional token scorer that is shared across all PTP layers to reduce
    parameter overhead. When `share_scorer=True` is passed to all PTP
    instances, they reference the same scorer module.
    """

    def __init__(self, dim: int, hidden: int = 64) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden),
            nn.GELU(),
            nn.Linear(hidden, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.sigmoid(self.net(x)).squeeze(-1)  # (B, L)

Optional token scorer that is shared across all PTP layers to reduce parameter overhead. When share_scorer=True is passed to all PTP instances, they reference the same scorer module.

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:
    return torch.sigmoid(self.net(x)).squeeze(-1)  # (B, L)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class SimpleNamespace (*args, **kwargs)

A simple attribute-based namespace.

class SpanExpandedAttention (config: Any,
rank: int | None = None,
expansion_span: int = 1024,
num_sink_tokens: int = 4,
use_lora: bool = True,
lora_rank: int = 8,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("span_expanded")
@ATTENTION_REGISTRY.register("span_expanded_attention")
class SpanExpandedAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any,
        rank: int | None = None,
        expansion_span: int = 1024,
        num_sink_tokens: int = 4,
        use_lora: bool = True,
        lora_rank: int = 8,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank or 64)
        self.config = config
        comp_config = getattr(config, "compression_config", None)

        self.hidden_size: int = getattr(config, "hidden_size", kwargs.get("hidden_size", 768))
        self.dim = self.hidden_size
        self.num_heads: int = getattr(
            config, "num_attention_heads", kwargs.get("num_attention_heads", 12)
        )
        self.num_kv_heads: int = getattr(
            config,
            "num_key_value_heads",
            kwargs.get("num_key_value_heads", self.num_heads),
        )
        self.head_dim: int = getattr(
            config,
            "head_dim",
            kwargs.get("head_dim", self.dim // self.num_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

        # ---- Span expansion configuration ----
        self.expansion_span: int = int(expansion_span) if expansion_span else 1024
        # Number of leading "sink" tokens to always preserve in the cache.
        # StreamingLLM (Xiao et al. 2024) shows the first few tokens act as
        # attention sinks and must be kept for long-context stability.
        self.num_sink_tokens: int = max(0, min(int(num_sink_tokens), self.expansion_span))

        is_dec = getattr(config, "is_decoder", kwargs.get("is_decoder", False))
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        self.lora_rank = (
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", lora_rank)),
            )
            if config
            else lora_rank
        )
        self.lora_alpha = (
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0
        self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.supports_gqa = True

        # Projections
        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(
                    config,
                    "hidden_dropout_prob",
                    getattr(config, "attention_probs_dropout_prob", 0.1),
                ),
            )
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if source_attn is not None:
            self.copy_weights(source_attn)

    # ------------------------------------------------------------------
    # Weight copying (unchanged, production-hardened)
    # ------------------------------------------------------------------
    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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (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])
                    ):
                        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. 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 (GPT-2 alt, Qwen, LLaMA, 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"SpanExpandedAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    else:
                        logger.warning(
                            "SpanExpandedAttention copy_weights: missing %s, using initialized weights",
                            proj_key,
                        )

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = torch.linalg.norm(target_mod.weight.float()).item()
                if fnorm < 1e-4:
                    logger.warning(
                        "SpanExpandedAttention copy_weights: %s has near-zero norm (%.2e)",
                        proj_key,
                        fnorm,
                    )

        if self.freeze_base:
            for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
                mod.freeze_base_weights()

    # ------------------------------------------------------------------
    # Cache management helpers
    # ------------------------------------------------------------------
    def _get_past_kv(
        self, past_key_value: Any
    ) -> tuple[torch.Tensor | None, torch.Tensor | None, int]:
        """Extract (past_k, past_v, past_len) from various cache formats."""
        if past_key_value is None:
            return None, None, 0

        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]
            return past_k, past_v, past_k.shape[-2]

        if hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
            past_k = past_key_value.key_cache[0]
            past_v = (
                past_key_value.value_cache[0] if hasattr(past_key_value, "value_cache") else None
            )
            return past_k, past_v, past_k.shape[-2]

        return None, None, 0

    def _truncate_cache_with_sinks(
        self, k: torch.Tensor, v: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Apply sliding-window truncation while preserving the first N sink tokens.

        StreamingLLM (Xiao et al., 2024): the first few tokens accumulate a
        disproportionate amount of attention and act as "attention sinks".
        Dropping them causes severe perplexity degradation in long contexts.
        We therefore always preserve the first `num_sink_tokens` entries.
        """
        span = self.expansion_span
        if span <= 0 or k.shape[-2] <= span:
            return k, v

        sink = self.num_sink_tokens
        if sink <= 0:
            # Pure rolling cache, no sinks preserved
            return k[..., -span:, :], v[..., -span:, :]

        if span <= sink:
            # Pathological config: keep only sinks, drop the rest
            return k[..., :span, :], v[..., :span, :]

        # Keep first `sink` tokens + most recent `span - sink` tokens
        sink_k = k[..., :sink, :]
        sink_v = v[..., :sink, :]
        recent_k = k[..., -(span - sink) :, :]
        recent_v = v[..., -(span - sink) :, :]
        return torch.cat([sink_k, recent_k], dim=-2), torch.cat([sink_v, recent_v], dim=-2)

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor, torch.Tensor] | None]:
        B, L, D = hidden_states.shape
        H, Dh = self.num_heads, self.head_dim

        # ---- QKV projections ----
        q = self.q_proj(hidden_states).view(B, L, H, Dh).permute(0, 2, 1, 3)
        k = self.k_proj(hidden_states).view(B, L, self.num_kv_heads, Dh).permute(0, 2, 1, 3)
        v = self.v_proj(hidden_states).view(B, L, self.num_kv_heads, Dh).permute(0, 2, 1, 3)

        # ---- Extract past cache ----
        past_k, past_v, past_len = self._get_past_kv(past_key_value)

        # ---- RoPE with past_len offset ----
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            cos_curr = (
                cos[..., past_len : past_len + L, :] if cos.shape[-2] >= past_len + L else cos
            )
            sin_curr = (
                sin[..., past_len : past_len + L, :] if sin.shape[-2] >= past_len + L else sin
            )
            q, k = apply_rotary_pos_emb(q, k, cos_curr, sin_curr)

        # ---- GQA repeat (before cache concat, so cache stays compact) ----
        # Note: If using GQA, we store the compact KV in cache and expand on use.
        # Here we expand first, which is simpler but uses more cache memory.
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            # Re-expand past cache if it was stored compact
            if past_k is not None and past_k.shape[1] != self.num_heads:
                past_k = repeat_kv(past_k, repeats)
                past_v = repeat_kv(past_v, repeats)
            k = repeat_kv(k, repeats)
            v = repeat_kv(v, repeats)

        # ---- Concatenate with past cache ----
        if past_k is not None:
            k = torch.cat([past_k, k], dim=-2)
            v = torch.cat([past_v, v], dim=-2)

        # ---- Sliding-window truncation with sink preservation ----
        k, v = self._truncate_cache_with_sinks(k, v)

        # ---- Present cache for next step ----
        # Store compact (un-repeated) KV if GQA is used to save memory:
        if self.num_heads != self.num_kv_heads and k.shape[1] == self.num_heads:
            repeats = self.num_heads // self.num_kv_heads
            # Take every `repeats`-th head to recover compact form
            # (repeat_kv duplicates contiguously, so index 0, repeats, 2*repeats, ...)
            compact_k = k[:, ::repeats, :, :].contiguous()
            compact_v = v[:, ::repeats, :, :].contiguous()
            present_key_value = (compact_k, compact_v)
        else:
            present_key_value = (k, v)

        scale = Dh**-0.5

        # ---- Fast path: PyTorch SDPA ----
        if attention_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 only meaningful when processing > 1 token at once.
                # When past cache is present and we process 1 token, no causal mask needed.
                _is_causal = bool(self.is_decoder) and L > 1 and past_len == 0
                _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, L, self.all_head_size)
                )
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                out = self.out_proj(_sdpa_out) if not self.is_nested_encoder else _sdpa_out
                return out, None, present_key_value
            except Exception:
                pass  # fall through to manual path

        # ---- Manual path ----
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        # Causal masking: query i (position past_len + i) attends to keys 0..(past_len + i)
        is_causal = self.is_decoder or kwargs.get("is_causal", False)
        if is_causal and L > 1:
            kv_len = k.shape[-2]
            causal_mask = torch.triu(
                torch.full((L, kv_len), -1e4, device=scores.device),
                diagonal=1 + past_len,
            )
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

        # External attention mask (padding etc.)
        if attention_mask is not None:
            scores = apply_attention_mask(scores, attention_mask)

        attn = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        # Head mask
        if head_mask is not None:
            if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                attn = attn * head_mask.view(1, self.num_heads, 1, 1).to(attn.dtype)
            elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                attn = attn * head_mask.view(head_mask.size(0), self.num_heads, 1, 1).to(attn.dtype)
            elif head_mask.dim() == 4:
                attn = attn * head_mask.to(attn.dtype)

        context = (
            torch.matmul(attn, v).permute(0, 2, 1, 3).contiguous().view(B, L, self.all_head_size)
        )

        if self.is_nested_encoder:
            return context, (attn if output_attentions else None), present_key_value

        out = self.out_proj(context)
        return out, (attn if output_attentions else None), present_key_value

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class SparsePruningCallback (model: Any,
prune_freq: int = 200,
final_sparsity: float = 0.7)

Expand source code

class SparsePruningCallback(_HFTrainerCallback):
    """
    Periodically hard-thresholds the sparse residual of every
    SpectronAttention module in the model.
    """

    def __init__(self, model: Any, prune_freq: int = 200, final_sparsity: float = 0.7):
        self.model = model
        self.prune_freq = prune_freq
        self.final_sparsity = final_sparsity

    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step > 0 and state.global_step % self.prune_freq == 0:
            for m in self.model.modules():
                if isinstance(m, SpectronAttention) and m.add_sparse:
                    m.prune_all_sparse(self.final_sparsity)

Periodically hard-thresholds the sparse residual of every SpectronAttention module in the model.

Ancestors

• transformers.trainer_callback.TrainerCallback

Methods

def on_step_end(self, args, state, control, **kwargs)

Expand source code

def on_step_end(self, args, state, control, **kwargs):
    if state.global_step > 0 and state.global_step % self.prune_freq == 0:
        for m in self.model.modules():
            if isinstance(m, SpectronAttention) and m.add_sparse:
                m.prune_all_sparse(self.final_sparsity)

Event called at the end of a training step. If using gradient accumulation, one training step might take several inputs.

class SpectraV2Attention (config: Any | None = None,
model_config: Any | None = None,
layer_idx: int = 0,
total_layers: int = 12,
bert_mode: bool = False,
pretrained_attn: nn.Module | None = None,
rank: int | None = 64,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("mha_sdpa")
@ATTENTION_REGISTRY.register("spectra_v2")
@ATTENTION_REGISTRY.register("spectra_v2_attention")
class SpectraV2Attention(BaseAttention):
    """
    MHA wrapper around PyTorch's scaled_dot_product_attention.
    Previously named spectra_v2 (deprecated alias).
    """

    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        model_config: Any | None = None,
        layer_idx: int = 0,
        total_layers: int = 12,
        bert_mode: bool = False,
        pretrained_attn: nn.Module | None = None,
        rank: int | None = 64,
        **kwargs: Any,
    ) -> None:
        cfg = config or model_config or kwargs.get("config")
        super().__init__(cfg, rank or 64)
        self.config = cfg
        comp_config = _cfg_get(cfg, "compression_config", kwargs.get("comp_config"))

        self.hidden_size = kwargs.get(
            "hidden_size",
            _cfg_get(cfg, "hidden_size", _cfg_get(cfg, "dim", _cfg_get(cfg, "d_model", 768))),
        )
        self.dim = self.hidden_size
        self.num_heads = kwargs.get(
            "num_heads",
            kwargs.get(
                "num_attention_heads",
                _cfg_get(
                    cfg,
                    "num_attention_heads",
                    _cfg_get(cfg, "num_heads", _cfg_get(cfg, "n_heads", 12)),
                ),
            ),
        )
        self.num_kv_heads = kwargs.get(
            "num_kv_heads",
            kwargs.get(
                "num_key_value_heads",
                _cfg_get(
                    cfg,
                    "num_key_value_heads",
                    _cfg_get(cfg, "num_kv_heads", self.num_heads),
                ),
            ),
        )
        self.head_dim = kwargs.get(
            "head_dim",
            _cfg_get(
                cfg,
                "head_dim",
                self.dim // self.num_heads if self.num_heads > 0 else 64,
            ),
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        self.layer_idx = layer_idx
        self.total_layers = total_layers
        self.bert_mode = bert_mode

        is_dec = _cfg_get(cfg, "is_decoder", kwargs.get("is_decoder", False))
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        self.lora_rank = _cfg_get(
            cfg,
            "lora_rank",
            _cfg_get(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha = _cfg_get(
            cfg,
            "lora_alpha",
            _cfg_get(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base = self.lora_rank > 0

        # Projections
        self.query = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.key = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.value = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        dropout_p = _cfg_get(cfg, "dropout", _cfg_get(cfg, "attention_probs_dropout_prob", 0.1))
        self.dropout = nn.Dropout(dropout_p)

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if source_attn is not None:
            self.copy_weights(source_attn)

    def copy_weights(self, src: Any) -> None:
        if src is None:
            return

        cls_name = src.__class__.__name__
        has_out_proj = (
            hasattr(src, "out_proj")
            or hasattr(src, "dense")
            or hasattr(src, "c_proj")
            or hasattr(src, "out_lin")
        )
        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_proj
        ):
            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 / c_proj
        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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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:
                extracted["o"] = _unpack_weight_bias(src.c_proj)

        # 2. DistilBERT style: q_lin, k_lin, v_lin, out_lin
        if hasattr(src, "q_lin") and src.q_lin is not None:
            extracted["q"] = _unpack_weight_bias(src.q_lin)
            extracted["k"] = _unpack_weight_bias(getattr(src, "k_lin", None))
            extracted["v"] = _unpack_weight_bias(getattr(src, "v_lin", None))
            extracted["o"] = _unpack_weight_bias(getattr(src, "out_lin", None))

        # 3. BERT / RoBERTa style: query, key, value, dense
        if hasattr(src, "query") and src.query is not None and extracted["q"][0] is None:
            extracted["q"] = _unpack_weight_bias(src.query)
            extracted["k"] = _unpack_weight_bias(getattr(src, "key", None))
            extracted["v"] = _unpack_weight_bias(getattr(src, "value", None))
            extracted["o"] = _unpack_weight_bias(
                getattr(src, "dense", getattr(src, "out_proj", None))
            )

        # 4. Nested structures (RoBERTa / BERT container objects)
        if hasattr(src, "self") and hasattr(src.self, "query") and extracted["q"][0] is None:
            s = src.self
            extracted["q"] = _unpack_weight_bias(getattr(s, "query", None))
            extracted["k"] = _unpack_weight_bias(getattr(s, "key", None))
            extracted["v"] = _unpack_weight_bias(getattr(s, "value", None))
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                extracted["o"] = _unpack_weight_bias(getattr(src.output, "dense", None))

        # 5. Qwen / LLaMA / HF Attention style: q_proj, k_proj, v_proj, o_proj
        if hasattr(src, "q_proj") and src.q_proj is not None and extracted["q"][0] is None:
            extracted["q"] = _unpack_weight_bias(src.q_proj)
            extracted["k"] = _unpack_weight_bias(getattr(src, "k_proj", None))
            extracted["v"] = _unpack_weight_bias(getattr(src, "v_proj", None))
            extracted["o"] = _unpack_weight_bias(
                getattr(src, "o_proj", getattr(src, "dense", getattr(src, "out_proj", None)))
            )

        # 6. UniversalAdapter fallback for custom names
        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)

        target_map = {
            "q": self.query,
            "k": self.key,
            "v": self.value,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    else:
                        logger.warning(
                            "Attention module copy_weights warning. Proceeding with initialized weights."
                        )
                        continue

                src_norm = torch.linalg.norm(w.float()).item()

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                elif (
                    w.shape[0] > target_mod.weight.shape[0]
                    and w.shape[0] % target_mod.weight.shape[0] == 0
                ):
                    repeats = w.shape[0] // target_mod.weight.shape[0]
                    # For GQA: average the query head groups
                    w_grouped = w.view(target_mod.weight.shape[0], repeats, w.shape[1]).mean(dim=1)
                    target_mod.weight.copy_(w_grouped)
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    elif (
                        b.shape[0] > target_mod.bias.shape[0]
                        and b.shape[0] % target_mod.bias.shape[0] == 0
                    ):
                        repeats = b.shape[0] // target_mod.bias.shape[0]
                        b_grouped = b.view(target_mod.bias.shape[0], repeats).mean(dim=1)
                        target_mod.bias.copy_(b_grouped)
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = torch.linalg.norm(target_mod.weight.float()).item()
                if src_norm > 1e-3 and fnorm < 1e-4:
                    raise RuntimeError(
                        f"SpectraV2Attention: copy_weights for {proj_key} produced near-zero norm "
                        f"({fnorm:.2e}) from a valid source ({src_norm:.2e}). Aborting to prevent silent "
                        f"model corruption."
                    )

        if self.freeze_base:
            for mod in [self.query, self.key, self.value, self.out_proj]:
                mod.freeze_base_weights()

    _copy_pretrained = copy_weights

    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> SpectraV2Attention:
        return cls(model_config=model_config, pretrained_attn=old_attn, **kwargs)

    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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, D = hidden_states.shape

        is_cross = encoder_hidden_states is not None
        kv_states = encoder_hidden_states if is_cross else hidden_states
        N_k = kv_states.shape[1]

        eff_mask = encoder_attention_mask if is_cross else attention_mask

        Q = self.query(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        K = self.key(kv_states).view(B, N_k, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
        V = self.value(kv_states).view(B, N_k, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)

        # Apply RoPE if provided
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # Repeat KV for GQA
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            K = repeat_kv(K, repeats)
            V = repeat_kv(V, repeats)

        scale = self.head_dim**-0.5

        # Causal determination
        _is_causal = getattr(self, "is_decoder", False) and not is_cross and N > 1

        attn_weights = None

        if 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

                attn_mask_for_sdpa = None
                if eff_mask is not None and not _is_causal:
                    if eff_mask.dim() == 2:
                        attn_mask_for_sdpa = eff_mask.bool().view(
                            eff_mask.shape[0], 1, 1, eff_mask.shape[1]
                        )
                    elif eff_mask.dim() == 4:
                        attn_mask_for_sdpa = eff_mask.bool()

                _sdpa_out = torch.nn.functional.scaled_dot_product_attention(
                    _qs,
                    _ks,
                    _vs,
                    attn_mask=attn_mask_for_sdpa,
                    is_causal=_is_causal,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )

                _sdpa_out = _sdpa_out.to(hidden_states.dtype)

                # Apply head_mask
                if head_mask is not None and head_mask.dim() >= 1:
                    hm = head_mask.view(-1, self.num_heads)
                    if hm.shape[0] == 1:
                        hm = hm.expand(_sdpa_out.shape[0], -1)
                    _sdpa_out = _sdpa_out * hm.view(_sdpa_out.shape[0], self.num_heads, 1, 1).to(
                        _sdpa_out.dtype
                    )

                _sdpa_out = _sdpa_out.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)

                context = _sdpa_out
                if self.is_nested_encoder:
                    if context.shape[-1] != hidden_states.shape[-1]:
                        raise RuntimeError(
                            f"SpectraV2Attention: output dimension mismatch. "
                            f"Got {context.shape[-1]}, expected {hidden_states.shape[-1]}."
                        )
                    return context, None

                output = self.out_proj(context)
                if output.shape[-1] != hidden_states.shape[-1]:
                    raise RuntimeError(
                        f"SpectraV2Attention: output dimension mismatch. "
                        f"Got {output.shape[-1]}, expected {hidden_states.shape[-1]}."
                    )
                return output, None

            except Exception as e:
                logger.debug("SDPA path failed or fell back: %s", e)
                pass

        # Fallback path (manual attention)
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if _is_causal:
            causal_mask = torch.triu(torch.full((N, N_k), -1e4, device=scores.device), diagonal=1)
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

        if eff_mask is not None:
            scores = apply_attention_mask(scores, eff_mask)

        attn_weights = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights)

        if head_mask is not None:
            if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                attn_weights = attn_weights * head_mask.view(1, self.num_heads, 1, 1).to(
                    attn_weights.dtype
                )
            elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                attn_weights = attn_weights * head_mask.view(
                    head_mask.size(0), self.num_heads, 1, 1
                ).to(attn_weights.dtype)
            elif head_mask.dim() == 4:
                attn_weights = attn_weights * head_mask.to(attn_weights.dtype)

        context = (
            torch.matmul(attn_weights, V)
            .permute(0, 2, 1, 3)
            .contiguous()
            .view(B, N, self.all_head_size)
        )

        if self.is_nested_encoder:
            if context.shape[-1] != hidden_states.shape[-1]:
                raise RuntimeError(
                    f"SpectraV2Attention: output dimension mismatch. "
                    f"Got {context.shape[-1]}, expected {hidden_states.shape[-1]}."
                )
            return context, (attn_weights if output_attentions else None)

        output = self.out_proj(context)
        if output.shape[-1] != hidden_states.shape[-1]:
            raise RuntimeError(
                f"SpectraV2Attention: output dimension mismatch. "
                f"Got {output.shape[-1]}, expected {hidden_states.shape[-1]}."
            )
        return output, (attn_weights if output_attentions else None)

MHA wrapper around PyTorch's scaled_dot_product_attention. Previously named spectra_v2 (deprecated alias).

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.

Static methods

def from_distilbert_pretrained(old_attn: nn.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> SpectraV2Attention

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class SpectraformerAttention (config: Any | None = None,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("spectraformer")
@ATTENTION_REGISTRY.register("spectraformer_attention")
class SpectraformerAttention(BaseAttention):
    """
    Spectraformer Attention: Exact linear-time attention via random feature kernelization.
    """

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank or 64)
        self.config = config
        comp_config = _cfg_get(config, "compression_config", None) or kwargs.get("comp_config")

        self.hidden_size = (
            kwargs.get("hidden_size")
            or _cfg_get(config, "hidden_size", None)
            or _cfg_get(config, "d_model", None)
        )
        if self.hidden_size is None and pretrained_attn is not None:
            for cand in ["q_lin", "query", "q_proj", "c_attn"]:
                mod = getattr(pretrained_attn, cand, None)
                if mod is not None and hasattr(mod, "weight"):
                    w = mod.weight
                    self.hidden_size = w.shape[1] if w.ndim == 2 else None
                    if self.hidden_size is not None:
                        break
        if self.hidden_size is None:
            self.hidden_size = 768
        self.dim = self.hidden_size

        self.num_heads = (
            kwargs.get("num_attention_heads")
            or kwargs.get("num_heads")
            or _cfg_get(config, "num_attention_heads", None)
            or _cfg_get(config, "num_heads", 12)
        )
        self.num_kv_heads = (
            kwargs.get("num_key_value_heads")
            or kwargs.get("num_kv_heads")
            or _cfg_get(
                config,
                "num_key_value_heads",
                self.num_heads,
            )
        )
        self.head_dim = _cfg_get(
            config,
            "head_dim",
            kwargs.get("head_dim", self.hidden_size // self.num_heads),
        )
        self.all_head_size = self.num_heads * self.head_dim
        self.all_kv_head_size = self.num_kv_heads * self.head_dim

        self.rank = _cfg_get(
            comp_config,
            "spectraformer_rank",
            _cfg_get(config, "spectraformer_rank", kwargs.pop("rank", rank or 64)),
        )

        # Causal decoder detection
        is_dec = _cfg_get(config, "is_decoder", kwargs.get("is_decoder", False))
        model_type = _cfg_get(config, "model_type", "").lower() if config else ""
        _mt = (model_type or "").lower()
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if any(_mt.startswith(p) for p in _decoder_prefixes):
            is_dec = True
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        # LoRA config
        self.lora_rank = _cfg_get(
            config,
            "lora_rank",
            _cfg_get(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha = _cfg_get(
            config,
            "lora_alpha",
            _cfg_get(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base = self.lora_rank > 0

        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=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.hidden_size,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.hidden_size,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        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=self.freeze_base,
        )

        # Aliases
        self.query = self.q_proj
        self.key = self.k_proj
        self.value = self.v_proj

        # Persistent orthogonal random feature matrix R: [head_dim, rank]
        # Generated once with QR decomposition for maximum feature variance coverage
        torch.manual_seed(42)
        H = torch.randn(self.head_dim, self.rank)
        Q_ortho, _ = torch.linalg.qr(H)
        self.register_buffer("R", Q_ortho.float())

        dropout_p = (
            getattr(config, "attention_probs_dropout_prob", kwargs.get("dropout", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if source_attn is not None:
            self.copy_weights(source_attn)

    def copy_weights(self, src: Any) -> None:
        """
        Robust weight copying across GPT-2, RoBERTa/BERT nested, Qwen/LLaMA, and DistilBERT.
        Validates Frobenius norm >= 1e-4.
        """
        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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (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])
                    ):
                        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. RoBERTa nested encoder
        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, Qwen, LLaMA)
        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"SpectraformerAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue

                src_norm = torch.linalg.norm(w.float()).item()
                if src_norm < 1e-4:
                    logger.warning(
                        f"Source weight for {proj_key} has near-zero norm {src_norm:.6f} < 1e-4"
                    )

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                    and w.shape[1] == target_mod.weight.shape[1]
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                fnorm = torch.linalg.norm(target_mod.weight.float()).item()
                if fnorm < 1e-4:
                    logger.warning(
                        "Compressor copy_weights warning. Proceeding with initialized weights."
                    )

        if self.freeze_base:
            for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
                mod.freeze_base_weights()

        # Preserve existing LoRA adapters from source projections
        for proj_key, target_mod in target_map.items():
            sub_name = {"q": "query", "k": "key", "v": "value", "o": "dense"}.get(proj_key, "")
            src_sub = getattr(src, sub_name, None)
            if src_sub is None and hasattr(src, f"{proj_key}_proj"):
                src_sub = getattr(src, f"{proj_key}_proj")
            if src_sub is None and hasattr(src, f"{proj_key}_lin"):
                src_sub = getattr(src, f"{proj_key}_lin")

            if src_sub is not None and hasattr(src_sub, "lora_A") and hasattr(src_sub, "lora_B"):
                if target_mod.lora_A is not None and hasattr(src_sub.lora_A, "weight"):
                    try:
                        target_mod.lora_A.data.copy_(src_sub.lora_A.weight.data)
                    except Exception as e:
                        logger.warning(f"copy_weights failed: {e}")
                if target_mod.lora_B is not None and hasattr(src_sub.lora_B, "weight"):
                    try:
                        target_mod.lora_B.data.copy_(src_sub.lora_B.weight.data)
                    except Exception as e:
                        logger.warning(f"copy_weights failed: {e}")

        return True

    def forward(
        self,
        hidden_states: torch.Tensor | None = None,
        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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        Forward pass with linear random feature kernel:
        1. Linear projection to Q, K, V
        2. Optional RoPE embedding
        3. GQA head expansion
        4. Random orthogonal feature projection with ReLU activation
        5. Linear attention context computation (recurrent for causal, global for bidirectional)
        """
        if hidden_states is None:
            hidden_states = kwargs.get("hidden_states") or (args[0] if args else None)
        if hidden_states is None:
            raise ValueError("SpectraformerAttention forward requires hidden_states")
        if attention_mask is None:
            attention_mask = kwargs.get("attention_mask")
        B, N, _ = hidden_states.shape

        q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # Apply RoPE if provided
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # GQA head expansion
        if self.num_heads != self.num_kv_heads:
            k = repeat_kv(k, self.num_heads // self.num_kv_heads)
            v = repeat_kv(v, self.num_heads // self.num_kv_heads)

        # Ensure R matches head_dim and rank
        R = self.R.to(dtype=hidden_states.dtype, device=hidden_states.device)
        if R.shape[0] != self.head_dim or R.shape[1] != self.rank:
            g = torch.Generator(
                device=(hidden_states.device if hidden_states.device.type != "cpu" else "cpu")
            ).manual_seed(42)
            H = torch.randn(
                self.head_dim,
                self.rank,
                device=hidden_states.device,
                dtype=torch.float32,
                generator=g,
            )
            R, _ = torch.linalg.qr(H)
            R = R.to(dtype=hidden_states.dtype)
            if hasattr(self, "R") and isinstance(self.R, torch.Tensor) and self.R.shape == R.shape:
                self.R.data.copy_(R)
            else:
                self.R = R

        scale = 1.0 / math.sqrt(self.head_dim)
        # Q_proj, K_proj: (B, H, N, r) - positive ELU+1 feature map with single scale on Q
        Q_proj = (F.elu(torch.matmul(q, R)) + 1.0) * scale
        K_proj = F.elu(torch.matmul(k, R)) + 1.0

        # Mask padding in K_proj
        if attention_mask is not None:
            valid_tok = get_token_valid_mask(attention_mask, B, N, hidden_states.device)
            v_mask_4d = valid_tok.unsqueeze(1).unsqueeze(-1).to(dtype=hidden_states.dtype)
            K_proj = K_proj * v_mask_4d
            v = v * v_mask_4d

        if self.is_decoder:
            # Recurrent linear causal attention: exact vectorized computation
            causal_mask = torch.tril(torch.ones(N, N, device=Q_proj.device, dtype=torch.bool))
            aff = torch.matmul(
                Q_proj.float(), K_proj.float().transpose(-2, -1).contiguous()
            )  # (B, H, N, N)
            aff = aff.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), 0.0)

            num = torch.matmul(aff, v.float())  # (B, H, N, d)
            den = aff.sum(dim=-1, keepdim=True) + 1e-8  # (B, H, N, 1)
            context = (num / den).to(hidden_states.dtype)
        else:
            # Bidirectional linear attention: exact O(N·r·d) computation
            KV = torch.matmul(
                K_proj.transpose(-2, -1).contiguous().float(), v.float()
            )  # (B, H, r, d)
            Z = (
                K_proj.float().sum(dim=2, keepdim=True).transpose(-2, -1).contiguous()
            )  # (B, H, r, 1)

            num = torch.matmul(Q_proj.float(), KV)  # (B, H, N, d)
            den = torch.matmul(Q_proj.float(), Z) + 1e-8  # (B, H, N, 1)
            context = (num / den).to(hidden_states.dtype)

        if head_mask is not None:
            hm = head_mask.to(context.device)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], 1, 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            context = context * hm

        context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)

        return (output, None)

Spectraformer Attention: Exact linear-time attention via random feature kernelization.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• BaseAttention
• torch.nn.modules.module.Module

Methods

def copy_weights(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    """
    Robust weight copying across GPT-2, RoBERTa/BERT nested, Qwen/LLaMA, and DistilBERT.
    Validates Frobenius norm >= 1e-4.
    """
    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()
            if (
                (c_attn_mod.__class__.__name__ == "Conv1D")
                or (base_c.__class__.__name__ == "Conv1D")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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()
                if (
                    (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])
                ):
                    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. RoBERTa nested encoder
    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, Qwen, LLaMA)
    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"SpectraformerAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}."
        )

    target_map = {
        "q": self.q_proj,
        "k": self.k_proj,
        "v": self.v_proj,
        "o": self.out_proj,
    }

    with torch.no_grad():
        for proj_key, target_mod in target_map.items():
            w, b = extracted[proj_key]
            if w is None:
                if proj_key == "o" and self.is_nested_encoder:
                    nn.init.eye_(target_mod.weight)
                    if target_mod.bias is not None:
                        nn.init.zeros_(target_mod.bias)
                    continue

            src_norm = torch.linalg.norm(w.float()).item()
            if src_norm < 1e-4:
                logger.warning(
                    f"Source weight for {proj_key} has near-zero norm {src_norm:.6f} < 1e-4"
                )

            if w.shape == target_mod.weight.shape:
                target_mod.weight.copy_(w)
            elif w.t().contiguous().shape == target_mod.weight.shape:
                target_mod.weight.copy_(w.t().contiguous())
            elif (
                w.shape[0] < target_mod.weight.shape[0]
                and target_mod.weight.shape[0] % w.shape[0] == 0
                and w.shape[1] == target_mod.weight.shape[1]
            ):
                repeats = target_mod.weight.shape[0] // w.shape[0]
                target_mod.weight.copy_(w.repeat(repeats, 1))
            else:
                min_0 = min(target_mod.weight.shape[0], w.shape[0])
                min_1 = min(target_mod.weight.shape[1], w.shape[1])
                target_mod.weight.zero_()
                target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

            if b is not None and target_mod.bias is not None:
                if b.shape == target_mod.bias.shape:
                    target_mod.bias.copy_(b)
                elif (
                    b.shape[0] < target_mod.bias.shape[0]
                    and target_mod.bias.shape[0] % b.shape[0] == 0
                ):
                    repeats = target_mod.bias.shape[0] // b.shape[0]
                    target_mod.bias.copy_(b.repeat(repeats))
                else:
                    min_b = min(target_mod.bias.shape[0], b.shape[0])
                    target_mod.bias.zero_()
                    target_mod.bias[:min_b].copy_(b[:min_b])
            elif target_mod.bias is not None:
                nn.init.zeros_(target_mod.bias)

            fnorm = torch.linalg.norm(target_mod.weight.float()).item()
            if fnorm < 1e-4:
                logger.warning(
                    "Compressor copy_weights warning. Proceeding with initialized weights."
                )

    if self.freeze_base:
        for mod in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
            mod.freeze_base_weights()

    # Preserve existing LoRA adapters from source projections
    for proj_key, target_mod in target_map.items():
        sub_name = {"q": "query", "k": "key", "v": "value", "o": "dense"}.get(proj_key, "")
        src_sub = getattr(src, sub_name, None)
        if src_sub is None and hasattr(src, f"{proj_key}_proj"):
            src_sub = getattr(src, f"{proj_key}_proj")
        if src_sub is None and hasattr(src, f"{proj_key}_lin"):
            src_sub = getattr(src, f"{proj_key}_lin")

        if src_sub is not None and hasattr(src_sub, "lora_A") and hasattr(src_sub, "lora_B"):
            if target_mod.lora_A is not None and hasattr(src_sub.lora_A, "weight"):
                try:
                    target_mod.lora_A.data.copy_(src_sub.lora_A.weight.data)
                except Exception as e:
                    logger.warning(f"copy_weights failed: {e}")
            if target_mod.lora_B is not None and hasattr(src_sub.lora_B, "weight"):
                try:
                    target_mod.lora_B.data.copy_(src_sub.lora_B.weight.data)
                except Exception as e:
                    logger.warning(f"copy_weights failed: {e}")

    return True

Robust weight copying across GPT-2, RoBERTa/BERT nested, Qwen/LLaMA, and DistilBERT. Validates Frobenius norm >= 1e-4.

def forward(self,
hidden_states: torch.Tensor | None = None,
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,
*args: Any,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor | None]

Expand source code

def forward(
    self,
    hidden_states: torch.Tensor | None = None,
    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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """
    Forward pass with linear random feature kernel:
    1. Linear projection to Q, K, V
    2. Optional RoPE embedding
    3. GQA head expansion
    4. Random orthogonal feature projection with ReLU activation
    5. Linear attention context computation (recurrent for causal, global for bidirectional)
    """
    if hidden_states is None:
        hidden_states = kwargs.get("hidden_states") or (args[0] if args else None)
    if hidden_states is None:
        raise ValueError("SpectraformerAttention forward requires hidden_states")
    if attention_mask is None:
        attention_mask = kwargs.get("attention_mask")
    B, N, _ = hidden_states.shape

    q = (
        self.q_proj(hidden_states)
        .view(B, N, self.num_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    k = (
        self.k_proj(hidden_states)
        .view(B, N, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v = (
        self.v_proj(hidden_states)
        .view(B, N, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # Apply RoPE if provided
    if position_embeddings is None:
        position_embeddings = kwargs.get("position_embeddings")
    if position_embeddings is not None:
        cos, sin = position_embeddings
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

    # GQA head expansion
    if self.num_heads != self.num_kv_heads:
        k = repeat_kv(k, self.num_heads // self.num_kv_heads)
        v = repeat_kv(v, self.num_heads // self.num_kv_heads)

    # Ensure R matches head_dim and rank
    R = self.R.to(dtype=hidden_states.dtype, device=hidden_states.device)
    if R.shape[0] != self.head_dim or R.shape[1] != self.rank:
        g = torch.Generator(
            device=(hidden_states.device if hidden_states.device.type != "cpu" else "cpu")
        ).manual_seed(42)
        H = torch.randn(
            self.head_dim,
            self.rank,
            device=hidden_states.device,
            dtype=torch.float32,
            generator=g,
        )
        R, _ = torch.linalg.qr(H)
        R = R.to(dtype=hidden_states.dtype)
        if hasattr(self, "R") and isinstance(self.R, torch.Tensor) and self.R.shape == R.shape:
            self.R.data.copy_(R)
        else:
            self.R = R

    scale = 1.0 / math.sqrt(self.head_dim)
    # Q_proj, K_proj: (B, H, N, r) - positive ELU+1 feature map with single scale on Q
    Q_proj = (F.elu(torch.matmul(q, R)) + 1.0) * scale
    K_proj = F.elu(torch.matmul(k, R)) + 1.0

    # Mask padding in K_proj
    if attention_mask is not None:
        valid_tok = get_token_valid_mask(attention_mask, B, N, hidden_states.device)
        v_mask_4d = valid_tok.unsqueeze(1).unsqueeze(-1).to(dtype=hidden_states.dtype)
        K_proj = K_proj * v_mask_4d
        v = v * v_mask_4d

    if self.is_decoder:
        # Recurrent linear causal attention: exact vectorized computation
        causal_mask = torch.tril(torch.ones(N, N, device=Q_proj.device, dtype=torch.bool))
        aff = torch.matmul(
            Q_proj.float(), K_proj.float().transpose(-2, -1).contiguous()
        )  # (B, H, N, N)
        aff = aff.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), 0.0)

        num = torch.matmul(aff, v.float())  # (B, H, N, d)
        den = aff.sum(dim=-1, keepdim=True) + 1e-8  # (B, H, N, 1)
        context = (num / den).to(hidden_states.dtype)
    else:
        # Bidirectional linear attention: exact O(N·r·d) computation
        KV = torch.matmul(
            K_proj.transpose(-2, -1).contiguous().float(), v.float()
        )  # (B, H, r, d)
        Z = (
            K_proj.float().sum(dim=2, keepdim=True).transpose(-2, -1).contiguous()
        )  # (B, H, r, 1)

        num = torch.matmul(Q_proj.float(), KV)  # (B, H, N, d)
        den = torch.matmul(Q_proj.float(), Z) + 1e-8  # (B, H, N, 1)
        context = (num / den).to(hidden_states.dtype)

    if head_mask is not None:
        hm = head_mask.to(context.device)
        if hm.dim() == 1:
            hm = hm.view(1, -1, 1, 1)
        elif hm.dim() == 2:
            hm = hm.view(hm.shape[0], 1, 1, 1)
        elif hm.dim() == 3:
            hm = hm.unsqueeze(1)
        context = context * hm

    context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size)

    if self.is_nested_encoder:
        output = context
    else:
        output = self.out_proj(context)

    return (output, None)

Forward pass with linear random feature kernel: 1. Linear projection to Q, K, V 2. Optional RoPE embedding 3. GQA head expansion 4. Random orthogonal feature projection with ReLU activation 5. Linear attention context computation (recurrent for causal, global for bidirectional)

Inherited members

• BaseAttention:
  • extra_repr

class SpectraformerRectify (head_dim: int, rff_features: int = 128, **kwargs)

Expand source code

class SpectraformerRectify(nn.Module):
    def __init__(self, head_dim: int, rff_features: int = 128, **kwargs) -> None:
        super().__init__()
        self.rff_features = rff_features
        proj_init = torch.randn(head_dim, rff_features)
        if head_dim <= rff_features:
            proj_q, _ = torch.linalg.qr(proj_init.T)
            proj_init = proj_q.T
        else:
            proj_q, _ = torch.linalg.qr(proj_init)
            proj_init = proj_q
        self.proj = nn.Parameter(proj_init / math.sqrt(head_dim))

    def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
        B, H, N, D = Q.shape
        F_dim = self.rff_features

        orig_dtype = Q.dtype
        Q, K, V = Q.float(), K.float(), V.float()
        proj = self.proj.to(Q.device).float()

        Q_proj = torch.matmul(F.normalize(Q, p=2, dim=-1), proj).clamp(max=11.0)
        K_proj = torch.matmul(F.normalize(K, p=2, dim=-1), proj).clamp(max=11.0)

        Q_feat = torch.exp(Q_proj) / math.sqrt(F_dim)
        K_feat = torch.exp(K_proj) / math.sqrt(F_dim)

        KV = torch.einsum("bhnf,bhnd->bhfd", K_feat, V)
        Z = K_feat.sum(dim=2)
        ctx = torch.einsum("bhnf,bhfd->bhnd", Q_feat, KV)
        norm = torch.einsum("bhnf,bhf->bhn", Q_feat, Z).unsqueeze(-1)
        # NX-03: clamp norm, not (norm + eps)
        return (ctx / norm.clamp(min=1e-5)).to(orig_dtype)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
    B, H, N, D = Q.shape
    F_dim = self.rff_features

    orig_dtype = Q.dtype
    Q, K, V = Q.float(), K.float(), V.float()
    proj = self.proj.to(Q.device).float()

    Q_proj = torch.matmul(F.normalize(Q, p=2, dim=-1), proj).clamp(max=11.0)
    K_proj = torch.matmul(F.normalize(K, p=2, dim=-1), proj).clamp(max=11.0)

    Q_feat = torch.exp(Q_proj) / math.sqrt(F_dim)
    K_feat = torch.exp(K_proj) / math.sqrt(F_dim)

    KV = torch.einsum("bhnf,bhnd->bhfd", K_feat, V)
    Z = K_feat.sum(dim=2)
    ctx = torch.einsum("bhnf,bhfd->bhnd", Q_feat, KV)
    norm = torch.einsum("bhnf,bhf->bhn", Q_feat, Z).unsqueeze(-1)
    # NX-03: clamp norm, not (norm + eps)
    return (ctx / norm.clamp(min=1e-5)).to(orig_dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class SpectralMamba (d_model: int = 768,
d_state: int = 16,
**kwargs: Any)

Expand source code

class SpectralMamba(nn.Module):
    """
    Spectral modulation with actual frequency filtering and low-rank state projection.

    Uses a learned frequency mask over FFT bins followed by low-rank adaptation.
    """

    def __init__(self, d_model: int = 768, d_state: int = 16, **kwargs: Any) -> None:
        super().__init__()
        self.d_model = d_model
        # Use a fixed resolution for frequency weights and interpolate to sequence length
        self.register_buffer(
            "freq_weights",
            torch.linspace(0.1, 1.0, 256).unsqueeze(1).expand(256, d_model).clone(),
        )
        self.register_buffer("dt_scale", torch.zeros(d_model))
        # Low-rank adaptation projection
        self.lora_down = nn.Parameter(torch.randn(d_model, d_state) * 0.01)
        self.lora_up = nn.Parameter(torch.zeros(d_state, d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Spectral modulation with actual frequency filtering.
        """
        B, L, D = x.shape
        dtype = x.dtype

        # FFT along sequence axis
        x_float = x.float()
        x_fft = torch.fft.rfft(x_float, dim=1)  # (B, L//2+1, D)

        freq_weights = self.freq_weights
        if freq_weights.shape[0] != x_fft.shape[1]:
            freq_weights = (
                F.interpolate(
                    freq_weights.transpose(0, 1).unsqueeze(0),
                    size=x_fft.shape[1],
                    mode="linear",
                    align_corners=False,
                )
                .squeeze(0)
                .transpose(0, 1)
            )

        x_fft_filtered = x_fft * freq_weights.unsqueeze(0)

        # Back to time domain
        x_filtered = torch.fft.irfft(x_fft_filtered, n=L, dim=1).to(dtype)

        # Low-rank state projection
        state_act = F.linear(
            F.linear(x_filtered, self.lora_down.t()),
            self.lora_up.t(),
        )

        # Learned channel gate
        gate = torch.sigmoid(self.dt_scale).unsqueeze(0).unsqueeze(0)

        return (x_filtered + state_act) * gate

Spectral modulation with actual frequency filtering and low-rank state projection.

Uses a learned frequency mask over FFT bins followed by low-rank adaptation.

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:
    """
    Spectral modulation with actual frequency filtering.
    """
    B, L, D = x.shape
    dtype = x.dtype

    # FFT along sequence axis
    x_float = x.float()
    x_fft = torch.fft.rfft(x_float, dim=1)  # (B, L//2+1, D)

    freq_weights = self.freq_weights
    if freq_weights.shape[0] != x_fft.shape[1]:
        freq_weights = (
            F.interpolate(
                freq_weights.transpose(0, 1).unsqueeze(0),
                size=x_fft.shape[1],
                mode="linear",
                align_corners=False,
            )
            .squeeze(0)
            .transpose(0, 1)
        )

    x_fft_filtered = x_fft * freq_weights.unsqueeze(0)

    # Back to time domain
    x_filtered = torch.fft.irfft(x_fft_filtered, n=L, dim=1).to(dtype)

    # Low-rank state projection
    state_act = F.linear(
        F.linear(x_filtered, self.lora_down.t()),
        self.lora_up.t(),
    )

    # Learned channel gate
    gate = torch.sigmoid(self.dt_scale).unsqueeze(0).unsqueeze(0)

    return (x_filtered + state_act) * gate

Spectral modulation with actual frequency filtering.

class SpectralMambaAttention (config: Any | None = None,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("spectral_mamba")
class SpectralMambaAttention(BaseAttention):
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank or 64)
        self.config = config
        comp_config = _cfg_get(config, "compression_config", kwargs.get("comp_config"))

        self.hidden_size: int = kwargs.get(
            "hidden_size",
            _cfg_get(
                config,
                "hidden_size",
                _cfg_get(config, "dim", _cfg_get(config, "d_model", 768)),
            ),
        )
        self.dim = self.hidden_size
        self.num_heads: int = kwargs.get(
            "num_heads",
            kwargs.get(
                "num_attention_heads",
                _cfg_get(
                    config,
                    "num_attention_heads",
                    _cfg_get(config, "num_heads", _cfg_get(config, "n_heads", 12)),
                ),
            ),
        )
        self.num_kv_heads: int = kwargs.get(
            "num_kv_heads",
            kwargs.get(
                "num_key_value_heads",
                _cfg_get(
                    config,
                    "num_key_value_heads",
                    _cfg_get(config, "num_kv_heads", self.num_heads),
                ),
            ),
        )
        self.head_dim: int = kwargs.get(
            "head_dim",
            _cfg_get(
                config,
                "head_dim",
                self.dim // self.num_heads if self.num_heads > 0 else 64,
            ),
        )
        self.all_head_size: int = self.num_heads * self.head_dim
        self.all_kv_head_size: int = self.num_kv_heads * self.head_dim

        is_dec = _cfg_get(config, "is_decoder", kwargs.get("is_decoder", False))
        model_type = str(_cfg_get(config, "model_type", "") or "").lower()
        decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if (
            any(model_type.startswith(p) for p in decoder_prefixes)
            or hasattr(config, "architectures")
            and any(
                "GPT" in str(a) or "CausalLM" in str(a)
                for a in (getattr(config, "architectures", []) or [])
            )
        ):
            is_dec = True
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = bool(is_dec)
        self.is_causal = bool(is_dec)
        self.is_nested_encoder = False

        self.lora_rank = _cfg_get(
            config,
            "lora_rank",
            _cfg_get(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
        )
        self.lora_alpha = _cfg_get(
            config,
            "lora_alpha",
            _cfg_get(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
        )
        self.freeze_base = self.lora_rank > 0

        # Projections
        self.q_lin = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.k_lin = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_lin = LoRALinear(
            self.dim,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.out_lin = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )

        self.q_proj = self.q_lin
        self.k_proj = self.k_lin
        self.v_proj = self.v_lin
        self.out_proj = self.out_lin
        self.query = self.q_lin
        self.key = self.k_lin
        self.value = self.v_lin

        d_state = _cfg_get(comp_config, "d_state", _cfg_get(config, "d_state", 16))
        self.spectral_mamba = SpectralMamba(self.hidden_size, d_state=d_state)
        # Gate implemented with frozen buffer weight/bias + LoRA adaptation so zero base parameters are added
        # Gate initialization:
        # output = gate * attn + (1 - gate) * mamba
        # With gate_bias=10.0, sigmoid(10.0) ≈ 0.99995, so the model starts as ~pure attention (identical to pretrained).
        self.register_buffer("gate_weight", torch.zeros(1, self.hidden_size))
        self.register_buffer("gate_bias", torch.tensor([10.0]))
        self.gate_lora_A = nn.Parameter(torch.randn(4, self.hidden_size) * 0.01)
        self.gate_lora_B = nn.Parameter(torch.zeros(1, 4))

        dropout_p = _cfg_get(
            config, "dropout", _cfg_get(config, "attention_probs_dropout_prob", 0.1)
        )
        self.dropout = nn.Dropout(dropout_p)

        source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if source_attn is not None:
            self.copy_weights(source_attn)

    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
        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()
                if (
                    (c_attn_mod.__class__.__name__ == "Conv1D")
                    or (base_c.__class__.__name__ == "Conv1D")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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()
                    if (
                        (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])
                    ):
                        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. 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)
                w, b = _unpack_weight_bias(mod)
                if w is not None:
                    extracted[proj_key] = (w, b)
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                out_mod = src.output.dense
                w, b = _unpack_weight_bias(out_mod)
                if w is not None:
                    extracted["o"] = (w, b)

        # 3. Flat modules
        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)
                        w, b = _unpack_weight_bias(mod)
                        if w is not None:
                            extracted[proj_key] = (w, b)
                            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:
            logger.warning(
                f"SpectralMambaAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}."
            )
            return

        target_map = {
            "q": self.q_lin,
            "k": self.k_lin,
            "v": self.v_lin,
            "o": self.out_lin,
        }

        with torch.no_grad():
            for proj_key, target_mod in target_map.items():
                w, b = extracted[proj_key]
                if w is None:
                    if proj_key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target_mod.weight)
                        if target_mod.bias is not None:
                            nn.init.zeros_(target_mod.bias)
                        continue
                    else:
                        logger.warning(f"Missing required weights for {proj_key}")
                        return

                if w.shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w)
                elif w.t().contiguous().shape == target_mod.weight.shape:
                    target_mod.weight.copy_(w.t().contiguous())
                elif (
                    w.shape[0] < target_mod.weight.shape[0]
                    and target_mod.weight.shape[0] % w.shape[0] == 0
                ):
                    repeats = target_mod.weight.shape[0] // w.shape[0]
                    # For GQA expansion or repeating
                    target_mod.weight.copy_(w.repeat(repeats, 1))
                elif (
                    w.shape[0] > target_mod.weight.shape[0]
                    and w.shape[0] % target_mod.weight.shape[0] == 0
                ):
                    repeats = w.shape[0] // target_mod.weight.shape[0]
                    # For GQA: mean-average the head groups from source
                    w_grouped = w.view(target_mod.weight.shape[0], repeats, w.shape[1]).mean(dim=1)
                    target_mod.weight.copy_(w_grouped)
                else:
                    min_0 = min(target_mod.weight.shape[0], w.shape[0])
                    min_1 = min(target_mod.weight.shape[1], w.shape[1])
                    target_mod.weight.zero_()
                    target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1])

                if b is not None and target_mod.bias is not None:
                    if b.shape == target_mod.bias.shape:
                        target_mod.bias.copy_(b)
                    elif (
                        b.shape[0] < target_mod.bias.shape[0]
                        and target_mod.bias.shape[0] % b.shape[0] == 0
                    ):
                        repeats = target_mod.bias.shape[0] // b.shape[0]
                        target_mod.bias.copy_(b.repeat(repeats))
                    elif (
                        b.shape[0] > target_mod.bias.shape[0]
                        and b.shape[0] % target_mod.bias.shape[0] == 0
                    ):
                        repeats = b.shape[0] // target_mod.bias.shape[0]
                        b_grouped = b.view(target_mod.bias.shape[0], repeats).mean(dim=1)
                        target_mod.bias.copy_(b_grouped)
                    else:
                        min_b = min(target_mod.bias.shape[0], b.shape[0])
                        target_mod.bias.zero_()
                        target_mod.bias[:min_b].copy_(b[:min_b])
                elif target_mod.bias is not None:
                    nn.init.zeros_(target_mod.bias)

                source_norm = torch.linalg.norm(w.float()).item()
                fnorm = torch.linalg.norm(target_mod.weight.float()).item()
                if fnorm < 1e-4:
                    if source_norm > 1e-3:
                        raise RuntimeError(
                            f"SpectralMambaAttention: weight copy for {proj_key} produced near-zero norm "
                            f"({fnorm:.2e}) from source norm ({source_norm:.2e}). Aborting."
                        )
                    logger.warning(f"Weight copy notice for {proj_key}: norm {fnorm:.6f} < 1e-4")

        if self.freeze_base:
            for mod in [self.q_lin, self.k_lin, self.v_lin, self.out_lin]:
                mod.freeze_base_weights()

    _copy_pretrained = copy_weights

    @classmethod
    def from_distilbert_pretrained(
        cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any
    ) -> SpectralMambaAttention:
        return cls(config=model_config, pretrained_attn=old_attn, **kwargs)

    def forward(
        self,
        hidden_states: torch.Tensor | None = None,
        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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if hidden_states is None:
            if "hidden_states" in kwargs and kwargs["hidden_states"] is not None:
                hidden_states = kwargs["hidden_states"]
            elif "x" in kwargs and kwargs["x"] is not None:
                hidden_states = kwargs["x"]
            elif "input_ids" in kwargs and kwargs["input_ids"] is not None:
                hidden_states = kwargs["input_ids"]
            elif len(args) > 0 and args[0] is not None:
                hidden_states = args[0]
        if hidden_states is None:
            raise ValueError(
                "SpectralMambaAttention forward requires hidden_states "
                "(pass as positional or via kwargs['hidden_states'])"
            )
        if attention_mask is None:
            attention_mask = kwargs.get("attention_mask")
        if hidden_states.dim() == 2:
            hidden_states = hidden_states.unsqueeze(0)
        B, N, D = hidden_states.shape

        q = self.q_lin(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k = (
            self.k_lin(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )
        v = (
            self.v_lin(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .permute(0, 2, 1, 3)
        )

        # Apply RoPE if provided
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # Repeat KV for GQA
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = repeat_kv(k, repeats)
            v = repeat_kv(v, repeats)

        scale = 1.0 / math.sqrt(self.head_dim)
        # Memory-efficient: use SDPA when attention weights are not requested
        attn_out = None
        attn_weights = None

        if not output_attentions and head_mask is None:
            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
                _sdpa_out = torch.nn.functional.scaled_dot_product_attention(
                    _qs,
                    _ks,
                    _vs,
                    is_causal=getattr(self, "is_decoder", False) and N > 1,
                    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, self.all_head_size)
                )
                _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4)
                attn_out = (
                    self.out_proj(_sdpa_out)
                    if not getattr(self, "is_nested_encoder", False)
                    else _sdpa_out
                )
            except Exception as e:
                logger.debug("SDPA failed, falling back to manual attention: %s", e)
                attn_out = None

        if attn_out is None:
            scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

            # Causal masking
            is_causal = self.is_decoder or kwargs.get("is_causal", False)
            if is_causal:
                causal_mask = torch.triu(torch.full((N, N), -1e4, device=scores.device), diagonal=1)
                scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

            # Attention mask
            if attention_mask is not None:
                scores = apply_attention_mask(scores, attention_mask)

            attn_weights = F.softmax(scores, dim=-1).to(hidden_states.dtype)
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            attn_weights = self.dropout(attn_weights)

            # head_mask handling
            if head_mask is not None:
                if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads:
                    attn_weights = attn_weights * head_mask.view(1, self.num_heads, 1, 1).to(
                        attn_weights.dtype
                    )
                elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads:
                    attn_weights = attn_weights * head_mask.view(
                        head_mask.size(0), self.num_heads, 1, 1
                    ).to(attn_weights.dtype)
                elif head_mask.dim() == 4:
                    attn_weights = attn_weights * head_mask.to(attn_weights.dtype)

            attn_ctx = (
                torch.matmul(attn_weights, v)
                .permute(0, 2, 1, 3)
                .contiguous()
                .view(B, N, self.all_head_size)
            )

            if self.is_nested_encoder:
                attn_out = attn_ctx
            else:
                attn_out = self.out_lin(attn_ctx)

        # Spectral Mamba stream
        mamba_out = self.spectral_mamba(hidden_states)
        while isinstance(mamba_out, (tuple, list)):
            mamba_out = mamba_out[0]

        if mamba_out.shape[-1] != attn_out.shape[-1]:
            if mamba_out.shape[-1] < attn_out.shape[-1]:
                mamba_out = F.pad(mamba_out, (0, attn_out.shape[-1] - mamba_out.shape[-1]))
            else:
                mamba_out = mamba_out[..., : attn_out.shape[-1]]

        # Adaptive gating: base linear + lora adapter
        gate_logits = F.linear(hidden_states, self.gate_weight, self.gate_bias)
        gate_logits = (
            gate_logits
            + F.linear(F.linear(hidden_states, self.gate_lora_A), self.gate_lora_B) * 2.0
        )
        g = torch.sigmoid(gate_logits)  # (B, N, 1)
        out = g * attn_out + (1.0 - g) * mamba_out

        if out.shape[-1] != hidden_states.shape[-1]:
            if out.shape[-1] < hidden_states.shape[-1]:
                out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1]))
            else:
                out = out[..., : hidden_states.shape[-1]]

        weights_to_return = None
        if output_attentions:
            weights_to_return = locals().get("attn_weights", None)
        return out, weights_to_return

Every 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.

Static methods

def from_distilbert_pretrained(old_attn: nn.Module,
model_config: Any | None = None,
**kwargs: Any) ‑> SpectralMambaAttention

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class SpectralMambaConfig (hidden_size: int = 768,
num_attention_heads: int = 12,
head_dim: int = 64,
d_state: int = 16,
d_conv: int = 4,
expand: int = 2,
dropout: float = 0.1,
gate_init_bias: float = 1.8)

Expand source code

@dataclass
class SpectralMambaConfig:
    hidden_size: int = 768
    num_attention_heads: int = 12
    head_dim: int = 64
    d_state: int = 16
    d_conv: int = 4
    expand: int = 2
    dropout: float = 0.1
    gate_init_bias: float = 1.8

SpectralMambaConfig(hidden_size: 'int' = 768, num_attention_heads: 'int' = 12, head_dim: 'int' = 64, d_state: 'int' = 16, d_conv: 'int' = 4, expand: 'int' = 2, dropout: 'float' = 0.1, gate_init_bias: 'float' = 1.8)

Instance variables

var d_conv : int

The type of the None singleton.

var d_state : int

The type of the None singleton.

var dropout : float

The type of the None singleton.

var expand : int

The type of the None singleton.

var gate_init_bias : float

The type of the None singleton.

var head_dim : int

The type of the None singleton.

var hidden_size : int

The type of the None singleton.

var num_attention_heads : int

The type of the None singleton.

class SpectronAttention (config: Any | None = None,
old_attn: Any | None = None,
dropout: float = 0.1,
layer_idx: int = 0,
ranks: dict[str, int] | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
add_sparse: bool = False,
sparsity_level: float = 0.7,
whitening: dict[str, torch.Tensor] | None = None,
seed: int = 42,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("spectron", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("spectron_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("spectron_attn", supports_causal=True, supports_gqa=True)
class SpectronAttention(BaseAttention):
    supports_causal: bool = True
    supports_gqa: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        old_attn: Any | None = None,
        dropout: float = 0.1,
        layer_idx: int = 0,
        ranks: dict[str, int] | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        add_sparse: bool = False,
        sparsity_level: float = 0.7,
        whitening: dict[str, torch.Tensor] | None = None,
        seed: int = 42,
        **kwargs: Any,
    ) -> None:
        model_config = config or kwargs.get("model_config")
        super().__init__(model_config)

        torch.manual_seed(seed)

        self.hidden_size = int(
            _cfg_get(model_config, "hidden_size", kwargs.get("hidden_size", 768))
        )
        self.num_heads = int(
            _cfg_get(
                model_config,
                "num_attention_heads",
                _cfg_get(model_config, "num_heads", kwargs.get("num_attention_heads", 12)),
            )
        )
        self.head_dim = int(
            _cfg_get(
                model_config,
                "head_dim",
                kwargs.get("head_dim", self.hidden_size // max(1, self.num_heads)),
            )
        )

        cfg_kv = _cfg_get(
            model_config,
            "num_key_value_heads",
            _cfg_get(model_config, "num_kv_heads", None),
        )
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.kv_repeat = max(1, self.num_heads // self.num_kv_heads)
        self.kv_head_size = self.num_kv_heads * self.head_dim
        self.q_head_size = self.num_heads * self.head_dim
        self.all_head_size = self.q_head_size

        # SP-FIX-04: explicit decoder detection
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        else:
            _mt = str(_cfg_get(model_config, "model_type", "") or "").lower()
            _dec_pfx = (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
            self.is_decoder = (
                bool(_cfg_get(model_config, "is_decoder", False))
                or bool(_cfg_get(model_config, "causal", False))
                or any(_mt.startswith(p) for p in _dec_pfx)
                or (old_attn is not None and hasattr(old_attn, "c_attn"))
            )

        self.is_nested_encoder = False

        # Ranks (with global clamp)
        default_r = int(_cfg_get(model_config, "svd_rank", _cfg_get(model_config, "rank", 64)))
        default_r = max(1, min(default_r, self.hidden_size, self.all_head_size))
        self.ranks: dict[str, int] = {}
        for k in ("q", "k", "v", "o"):
            val = ranks.get(k, default_r) if ranks else default_r
            self.ranks[k] = max(1, min(int(val), self.hidden_size, self.all_head_size))

        self.add_sparse = bool(add_sparse)
        self.sparsity_level = float(sparsity_level)
        self.lora_rank = max(0, int(lora_rank)) if lora_rank else 0
        self.lora_alpha = float(lora_alpha)
        self.scaling = (lora_alpha / self.lora_rank) if self.lora_rank > 0 else 1.0

        # Whitening factors per projection (from calibration)
        self.whitening = whitening or {}

        # ---- factorised base projections (SVD) ----
        self.q_down = nn.Linear(self.hidden_size, self.ranks["q"], bias=False)
        self.q_up = nn.Linear(self.ranks["q"], self.q_head_size, bias=True)
        self.k_down = nn.Linear(self.hidden_size, self.ranks["k"], bias=False)
        self.k_up = nn.Linear(self.ranks["k"], self.kv_head_size, bias=True)
        self.v_down = nn.Linear(self.hidden_size, self.ranks["v"], bias=False)
        self.v_up = nn.Linear(self.ranks["v"], self.kv_head_size, bias=True)
        self.o_down = nn.Linear(self.q_head_size, self.ranks["o"], bias=False)
        self.o_up = nn.Linear(self.ranks["o"], self.hidden_size, bias=True)

        for mod in (
            self.q_down,
            self.q_up,
            self.k_down,
            self.k_up,
            self.v_down,
            self.v_up,
            self.o_down,
            self.o_up,
        ):
            nn.init.orthogonal_(mod.weight)
            if mod.bias is not None:
                nn.init.zeros_(mod.bias)
            mod.weight.requires_grad = False
            if mod.bias is not None:
                mod.bias.requires_grad = False

        # ---- LoRA adapters ----
        if self.lora_rank > 0:
            self.q_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.q_lora_B = nn.Parameter(torch.zeros(self.q_head_size, self.lora_rank))
            self.k_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.k_lora_B = nn.Parameter(torch.zeros(self.kv_head_size, self.lora_rank))
            self.v_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.hidden_size))
            self.v_lora_B = nn.Parameter(torch.zeros(self.kv_head_size, self.lora_rank))
            self.o_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.q_head_size))
            self.o_lora_B = nn.Parameter(torch.zeros(self.hidden_size, self.lora_rank))
            for p_a in (self.q_lora_A, self.k_lora_A, self.v_lora_A, self.o_lora_A):
                nn.init.kaiming_uniform_(p_a, a=math.sqrt(5))
        else:
            self.register_parameter("q_lora_A", None)
            self.register_parameter("q_lora_B", None)
            self.register_parameter("k_lora_A", None)
            self.register_parameter("k_lora_B", None)
            self.register_parameter("v_lora_A", None)
            self.register_parameter("v_lora_B", None)
            self.register_parameter("o_lora_A", None)
            self.register_parameter("o_lora_B", None)

        # ---- optional sparse residuals (init from zero) ----
        if self.add_sparse:
            self.q_sparse = nn.Parameter(torch.zeros(self.q_head_size, self.hidden_size))
            self.k_sparse = nn.Parameter(torch.zeros(self.kv_head_size, self.hidden_size))
            self.v_sparse = nn.Parameter(torch.zeros(self.kv_head_size, self.hidden_size))
            self.o_sparse = nn.Parameter(torch.zeros(self.hidden_size, self.q_head_size))
        else:
            self.register_parameter("q_sparse", None)
            self.register_parameter("k_sparse", None)
            self.register_parameter("v_sparse", None)
            self.register_parameter("o_sparse", None)

        dropout_p = _cfg_get(
            model_config,
            "dropout",
            _cfg_get(
                model_config,
                "hidden_dropout_prob",
                _cfg_get(model_config, "attention_probs_dropout_prob", dropout),
            ),
        )
        self.dropout = nn.Dropout(float(dropout_p) if dropout_p is not None else 0.1)

        src_attn = old_attn or kwargs.get("pretrained_attn")
        if src_attn is not None:
            self.copy_weights(src_attn)

    # ------------------------------------------------------------------
    # SP-FIX-01: activation-aware factorization
    # ------------------------------------------------------------------
    def _factorise_and_set(
        self,
        W: torch.Tensor,
        bias: torch.Tensor | None,
        down: nn.Linear,
        up: nn.Linear,
        rank: int,
        whitening_L: torch.Tensor | None = None,
    ) -> None:
        """
        Factories W into down/up via activation-aware SVD:

            if whitening_L is None:
                U Σ V^T = SVD(W)
                down = √Σ V^T
                up   = U √Σ
            else:
                U Σ V^T = SVD(W L)
                down = √Σ V^T L^{-1}
                up   = U √Σ
        """
        with torch.no_grad():
            W_f = W.float()
            d_out, d_in = W_f.shape
            r = min(rank, d_out, d_in, down.out_features, up.in_features)
            r = max(r, 1)

            L = None
            if whitening_L is not None and whitening_L.shape == (d_in, d_in):
                L = whitening_L.float().to(W_f.device)

            if L is not None:
                eps = 1e-6 * torch.trace(L @ L.t()).item() / max(1, d_in)
                eye = torch.eye(d_in, device=W_f.device, dtype=W_f.dtype)
                try:
                    L_reg = L + eps * eye  # small regularization
                    W_eff = W_f @ L_reg
                    U, s, Vt = torch.linalg.svd(W_eff, full_matrices=False)
                    L_inv = torch.linalg.inv(L_reg)
                    Vt = Vt @ L_inv  # bring back to original coordinate
                except Exception as e:
                    log.debug("Whitening failed on factorise (%s); fallback.", e)
                    U, s, Vt = torch.linalg.svd(W_f, full_matrices=False)
            else:
                U, s, Vt = torch.linalg.svd(W_f, full_matrices=False)

            sqrt_s = torch.sqrt(torch.clamp(s[:r], min=1e-12))
            down_w = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(down.weight.device, down.weight.dtype)
            up_w = (U[:, :r] * sqrt_s.unsqueeze(0)).to(up.weight.device, up.weight.dtype)

            down.weight.data.zero_()
            up.weight.data.zero_()
            down.weight.data[:r, :].copy_(down_w)
            up.weight.data[:, :r].copy_(up_w)
            if bias is not None and up.bias is not None:
                up.bias.data.copy_(bias.to(up.bias.device, up.bias.dtype))

    # ------------------------------------------------------------------
    # copy_weights with activation-aware factorization
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_out_proj = any(
                hasattr(old_attn, a)
                for a in (
                    "dense",
                    "out_proj",
                    "c_proj",
                    "out_lin",
                    "o_proj",
                    "wo",
                    "o",
                    "proj",
                )
            )
            if src_cls_name in (
                "BertSelfAttention",
                "RobertaSelfAttention",
                "XLMRobertaSelfAttention",
                "ElectraSelfAttention",
                "DebertaSelfAttention",
            ) or (hasattr(old_attn, "query") and not has_out_proj):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o, b_o = _unpack_weight_bias(old_attn.output.dense)
            else:
                src = old_attn

            # GPT-2 fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_w, c_b = _unpack_weight_bias(src.c_attn)
                if c_w is not None:
                    if c_w.shape[0] == self.hidden_size and c_w.shape[1] != self.hidden_size:
                        c_w = c_w.t().contiguous()
                    q_dim = self.q_head_size
                    k_dim = self.hidden_size
                    v_dim = self.hidden_size
                    w_q = c_w[:q_dim, :]
                    w_k = c_w[q_dim : q_dim + k_dim, :]
                    w_v = c_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
                    if c_b is not None:
                        b_q = c_b[:q_dim]
                        b_k = c_b[q_dim : q_dim + k_dim]
                        b_v = c_b[q_dim + k_dim : q_dim + k_dim + v_dim]

                out_w, out_b = _unpack_weight_bias(src.c_proj)
                if out_w is not None:
                    if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                        out_w = out_w.t().contiguous()
                    w_o, b_o = out_w, out_b

            # BERT / RoBERTa nested
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q, b_q = _unpack_weight_bias(src.query)
                w_k, b_k = _unpack_weight_bias(src.key)
                w_v, b_v = _unpack_weight_bias(src.value)
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            w_o, b_o = _unpack_weight_bias(getattr(src, attr))
                            break

            # Qwen / LLaMA
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q, b_q = _unpack_weight_bias(src.q_proj)
                w_k, b_k = _unpack_weight_bias(src.k_proj)
                w_v, b_v = _unpack_weight_bias(src.v_proj)
                if hasattr(src, "o_proj"):
                    w_o, b_o = _unpack_weight_bias(src.o_proj)

            # DistilBERT
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q, b_q = _unpack_weight_bias(src.q_lin)
                w_k, b_k = _unpack_weight_bias(src.k_lin)
                w_v, b_v = _unpack_weight_bias(src.v_lin)
                if hasattr(src, "out_lin"):
                    w_o, b_o = _unpack_weight_bias(src.out_lin)

            # UniversalAdapter fallback
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter(old_attn)
                    extracted = adapter.extract_qkv_out()
                    if extracted:
                        w_q = extracted.get("q_weight")
                        b_q = extracted.get("q_bias")
                        w_k = extracted.get("k_weight")
                        b_k = extracted.get("k_bias")
                        w_v = extracted.get("v_weight")
                        b_v = extracted.get("v_bias")
                        if w_o is None:
                            w_o = extracted.get("out_weight")
                            b_o = extracted.get("out_bias")
                except Exception as e:
                    log.debug("UniversalAdapter fallback failed: %s", e)

            missing = []
            if w_q is None:
                missing.append("query")
            if w_k is None:
                missing.append("key")
            if w_v is None:
                missing.append("value")
            if missing:
                logger.warning("Spectron: missing %s; keeping random init.", missing)
                return

            def _orient(w, out_dim, in_dim):
                if w.shape == (out_dim, in_dim):
                    return w
                if w.shape == (in_dim, out_dim):
                    return w.t().contiguous()
                return w

            # ---- activation-aware factorization using self.whitening ----
            w_q = _orient(w_q, self.q_head_size, self.hidden_size)
            self._factorise_and_set(
                w_q,
                b_q,
                self.q_down,
                self.q_up,
                self.ranks["q"],
                whitening_L=self.whitening.get("q", None),
            )

            # K with GQA pooling
            if w_k.shape[0] == self.kv_head_size:
                w_k = _orient(w_k, self.kv_head_size, self.hidden_size)
                self._factorise_and_set(
                    w_k,
                    b_k,
                    self.k_down,
                    self.k_up,
                    self.ranks["k"],
                    whitening_L=self.whitening.get("k", None),
                )
            else:
                rep = max(1, w_k.shape[0] // self.kv_head_size)
                w_k_pooled = (
                    w_k.view(self.num_kv_heads, rep, self.head_dim, -1)
                    .mean(dim=1)
                    .view(self.kv_head_size, -1)
                )
                b_k_pooled = (
                    b_k.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                    if b_k is not None
                    else None
                )
                self._factorise_and_set(
                    w_k_pooled,
                    b_k_pooled,
                    self.k_down,
                    self.k_up,
                    self.ranks["k"],
                    whitening_L=self.whitening.get("k", None),
                )

            # V with GQA pooling
            if w_v.shape[0] == self.kv_head_size:
                w_v = _orient(w_v, self.kv_head_size, self.hidden_size)
                self._factorise_and_set(
                    w_v,
                    b_v,
                    self.v_down,
                    self.v_up,
                    self.ranks["v"],
                    whitening_L=self.whitening.get("v", None),
                )
            else:
                rep = max(1, w_v.shape[0] // self.kv_head_size)
                w_v_pooled = (
                    w_v.view(self.num_kv_heads, rep, self.head_dim, -1)
                    .mean(dim=1)
                    .view(self.kv_head_size, -1)
                )
                b_v_pooled = (
                    b_v.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1)
                    if b_v is not None
                    else None
                )
                self._factorise_and_set(
                    w_v_pooled,
                    b_v_pooled,
                    self.v_down,
                    self.v_up,
                    self.ranks["v"],
                    whitening_L=self.whitening.get("v", None),
                )

            # Output projection
            if w_o is not None and not self.is_nested_encoder:
                w_o = _orient(w_o, self.hidden_size, self.q_head_size)
                self._factorise_and_set(
                    w_o,
                    b_o,
                    self.o_down,
                    self.o_up,
                    self.ranks["o"],
                    whitening_L=self.whitening.get("o", None),
                )

    # ------------------------------------------------------------------
    # Forward helpers
    # ------------------------------------------------------------------
    def _linear_proj(
        self,
        x: torch.Tensor,
        down: nn.Linear,
        up: nn.Linear,
        lora_A: nn.Parameter | None,
        lora_B: nn.Parameter | None,
        sparse_w: nn.Parameter | None,
    ) -> torch.Tensor:
        out = up(down(x))
        if self.lora_rank > 0 and lora_A is not None and lora_B is not None:
            a = lora_A.to(x.dtype)
            b = lora_B.to(x.dtype)
            out = out + F.linear(F.linear(x, a), b) * self.scaling
        if self.add_sparse and sparse_w is not None:
            out = out + F.linear(x, sparse_w.to(x.dtype))
        return out

    # ------------------------------------------------------------------
    # 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,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        kv_len = kv_input.shape[1]

        Q = self._linear_proj(
            hidden_states,
            self.q_down,
            self.q_up,
            self.q_lora_A,
            self.q_lora_B,
            self.q_sparse,
        )
        K = self._linear_proj(
            kv_input,
            self.k_down,
            self.k_up,
            self.k_lora_A,
            self.k_lora_B,
            self.k_sparse,
        )
        V = self._linear_proj(
            kv_input,
            self.v_down,
            self.v_up,
            self.v_lora_A,
            self.v_lora_B,
            self.v_sparse,
        )

        Q = Q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        K = K.view(B, kv_len, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()
        V = V.view(B, kv_len, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous()

        # RoPE
        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 = 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]
            cos = cos[..., past_len : past_len + N, :]
            sin = sin[..., past_len : past_len + N, :]
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # GQA
        if self.num_kv_heads < self.num_heads:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        scale = 1.0 / math.sqrt(self.head_dim)

        # SP-FIX-04: correct causal condition
        is_causal = self.is_decoder and not is_cross and N > 1

        # SP-FIX-05: removed padding/truncation hack
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        if is_causal:
            causal = (
                torch.triu(
                    torch.full(
                        (N, kv_len),
                        float("-inf"),
                        device=scores.device,
                        dtype=torch.float32,
                    ),
                    diagonal=1 + max(0, kv_len - N),
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                scores = scores + (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
            elif eff_mask.dim() == 3:
                scores = scores + eff_mask.unsqueeze(1).float()
            elif eff_mask.dim() == 4:
                scores = scores + eff_mask.float()

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn_weights = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights)

        if head_mask is not None:
            hm = head_mask.to(attn_weights.dtype)
            if hm.dim() == 1:
                hm = hm.view(1, -1, 1, 1)
            elif hm.dim() == 2:
                hm = hm.view(hm.shape[0], hm.shape[1], 1, 1)
            elif hm.dim() == 3:
                hm = hm.unsqueeze(1)
            attn_weights = attn_weights * hm

        context = (
            torch.matmul(attn_weights, V)
            .transpose(1, 2)
            .contiguous()
            .view(B, N, self.all_head_size)
        )
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = context
        else:
            output = self._linear_proj(
                context,
                self.o_down,
                self.o_up,
                self.o_lora_A,
                self.o_lora_B,
                self.o_sparse,
            )

        return output, (attn_weights if output_attentions else None)

    # ------------------------------------------------------------------
    def total_sparse_l1(self) -> torch.Tensor:
        """Sum of L1 norms of all sparse residuals. Add to training loss."""
        if not self.add_sparse:
            return torch.zeros((), device=self.q_down.weight.device)
        total = torch.zeros((), device=self.q_down.weight.device)
        for s in (self.q_sparse, self.k_sparse, self.v_sparse, self.o_sparse):
            if s is not None:
                total = total + s.abs().sum()
        return total

    def prune_all_sparse(self, sparsity: float) -> None:
        with torch.no_grad():
            for s in (self.q_sparse, self.k_sparse, self.v_sparse, self.o_sparse):
                if s is None or s.numel() == 0:
                    continue
                keep = int(s.numel() * (1.0 - sparsity))
                if keep <= 0:
                    s.zero_()
                    continue
                if keep >= s.numel():
                    continue
                flat = s.abs().view(-1)
                kth = torch.kthvalue(flat, s.numel() - keep).values
                s.mul_(s.abs() >= kth)

Every 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 prune_all_sparse(self, sparsity: float) ‑> None

Expand source code

def prune_all_sparse(self, sparsity: float) -> None:
    with torch.no_grad():
        for s in (self.q_sparse, self.k_sparse, self.v_sparse, self.o_sparse):
            if s is None or s.numel() == 0:
                continue
            keep = int(s.numel() * (1.0 - sparsity))
            if keep <= 0:
                s.zero_()
                continue
            if keep >= s.numel():
                continue
            flat = s.abs().view(-1)
            kth = torch.kthvalue(flat, s.numel() - keep).values
            s.mul_(s.abs() >= kth)

def total_sparse_l1(self) ‑> torch.Tensor

Expand source code

def total_sparse_l1(self) -> torch.Tensor:
    """Sum of L1 norms of all sparse residuals. Add to training loss."""
    if not self.add_sparse:
        return torch.zeros((), device=self.q_down.weight.device)
    total = torch.zeros((), device=self.q_down.weight.device)
    for s in (self.q_sparse, self.k_sparse, self.v_sparse, self.o_sparse):
        if s is not None:
            total = total + s.abs().sum()
    return total

Sum of L1 norms of all sparse residuals. Add to training loss.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class SpectronSpectralEntropyRouter (config: AegisSpectronFP8Config)

Expand source code

class SpectronSpectralEntropyRouter(nn.Module):
    def __init__(self, config: AegisSpectronFP8Config):
        super().__init__()
        self.tau = float(config.temperature)
        self.lam = float(config.lambda_cost)
        self.register_buffer("costs", torch.tensor([config.c_loc, config.c_glob, config.c_lin]))
        self.h_proj = nn.Linear(config.hidden_size, 3)
        nn.init.zeros_(self.h_proj.weight)
        self.h_proj.bias.data.copy_(
            torch.tensor(config.router_init_bias, dtype=self.h_proj.bias.dtype)
        )

    def forward(self, X: torch.Tensor, token_valid: torch.Tensor | None = None) -> torch.Tensor:
        B, N, D = X.shape

        with torch.no_grad():
            S_sq = X.float().pow(2)
            denom = S_sq.sum(dim=-1, keepdim=True).clamp(min=1e-6)
            p = (S_sq / denom).clamp(min=1e-8, max=1.0)
            entropy = -(p * torch.log(p)).sum(dim=-1)  # (B, N)

            if token_valid is not None:
                entropy = entropy * token_valid.float()
                entropy = entropy / entropy.amax(dim=-1, keepdim=True).clamp(min=1e-6)

        h = self.h_proj(X.to(self.h_proj.weight.dtype))  # (B, N, 3)
        h = h * entropy.unsqueeze(-1).to(h.dtype)

        c = self.costs.to(h.dtype)
        logits = (h - self.lam * c) / max(self.tau, 1e-4)
        return F.softmax(logits.float(), dim=-1).to(X.dtype)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, X: torch.Tensor, token_valid: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def forward(self, X: torch.Tensor, token_valid: torch.Tensor | None = None) -> torch.Tensor:
    B, N, D = X.shape

    with torch.no_grad():
        S_sq = X.float().pow(2)
        denom = S_sq.sum(dim=-1, keepdim=True).clamp(min=1e-6)
        p = (S_sq / denom).clamp(min=1e-8, max=1.0)
        entropy = -(p * torch.log(p)).sum(dim=-1)  # (B, N)

        if token_valid is not None:
            entropy = entropy * token_valid.float()
            entropy = entropy / entropy.amax(dim=-1, keepdim=True).clamp(min=1e-6)

    h = self.h_proj(X.to(self.h_proj.weight.dtype))  # (B, N, 3)
    h = h * entropy.unsqueeze(-1).to(h.dtype)

    c = self.costs.to(h.dtype)
    logits = (h - self.lam * c) / max(self.tau, 1e-4)
    return F.softmax(logits.float(), dim=-1).to(X.dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class SqueezingHeads (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
head_keep_ratio: float | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("squeezing_heads_distiller", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("squeezing_heads_distill", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("squeezing_heads", supports_causal=True, supports_gqa=True)
class SqueezingHeadsDistiller(BaseAttention):
    """
    Attention wrapper that keeps only the top-k most important heads
    (based on pretrained Q/K/V Frobenius norms) while preserving GQA structure.

    Effective head count:
        num_kept_heads = round(head_keep_ratio * num_heads)
        q_head_size   = num_kept_heads * head_dim
        out_head_size = num_kept_heads * head_dim

    KV projections preserve all num_kv_heads (GQA). Kept Q heads are grouped
    to matching KV heads via GQA index mapping.

    Parameter reduction:
        Q : num_heads   * head_dim * hidden  →  num_kept_heads * head_dim * hidden
        O : hidden * num_heads * head_dim    →  hidden * num_kept_heads * head_dim
        K/V: unchanged (GQA-compatible)
    """

    supports_gqa: bool = True
    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = None,
        head_keep_ratio: float | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        # ---- 1. Dimensions ----
        self.hidden_size = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.num_heads = int(
            _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "num_heads", kwargs.get("num_attention_heads", 12)),
            )
        )

        cfg_hd = _cfg_get(config, "head_dim", None)
        kw_hd = kwargs.get("head_dim")
        if cfg_hd is not None:
            self.head_dim = int(cfg_hd)
        elif kw_hd is not None:
            self.head_dim = int(kw_hd)
        elif old_attn is not None:
            self.head_dim = None
            for attr in ("q_proj", "query", "q_lin"):
                if hasattr(old_attn, attr):
                    qw = getattr(old_attn, attr).weight
                    self.head_dim = qw.shape[0] // max(1, self.num_heads)
                    break
            if self.head_dim is None:
                self.head_dim = self.hidden_size // max(1, self.num_heads)
        else:
            self.head_dim = self.hidden_size // max(1, self.num_heads)

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))

        # ---- 2. Compute num_kept_heads ----
        # SQ2-FIX-01: derive head_keep_ratio from `rank` if not explicitly given.
        if head_keep_ratio is not None:
            hkr = float(head_keep_ratio)
        elif rank is not None and rank > 0:
            # If rank <= num_heads, interpret as "keep this many heads"
            if rank <= self.num_heads:
                hkr = float(rank) / float(self.num_heads)
            else:
                # Otherwise interpret as "target head_dim"
                effective_hd = min(int(rank), self.head_dim)
                hkr = float(effective_hd) / float(self.head_dim)
        else:
            hkr = 1.0
        self.head_keep_ratio = max(1.0 / self.num_heads, min(1.0, hkr))

        self.num_kept_heads = max(
            1, min(self.num_heads, int(round(self.num_heads * self.head_keep_ratio)))
        )
        self.head_keep_ratio = self.num_kept_heads / self.num_heads

        self.kept_q_head_size = self.num_kept_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim
        self.all_head_size = self.kept_q_head_size

        # GQA mapping: which KV head does each kept Q head map to?
        # Kept heads are the top-k by importance; we must remember their original indices.
        self.num_kv_groups = max(1, self.num_heads // self.num_kv_heads)
        self.kept_head_indices: torch.Tensor | None = None  # filled by copy_weights

        # ---- 3. Decoder detection ----
        self.is_decoder = _detect_decoder(config, is_decoder, old_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        self.lora_rank = int(lora_rank if lora_rank is not None else 16)
        self.lora_alpha = float(lora_alpha)

        # ---- 4. Projections ----
        # SQ2-FIX-01: Q and O are sized to num_kept_heads, not num_heads.
        self.q_proj = LoRALinear(
            self.hidden_size,
            self.kept_q_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.kept_q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        # ---- 5. Copy weights from source ----
        src_attn = old_attn if old_attn is not None else kwargs.get("pretrained_attn")
        if src_attn is not None:
            self.copy_weights(src_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 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

    # ------------------------------------------------------------------
    # Head importance: SQ2-FIX-02
    # ------------------------------------------------------------------
    def _compute_head_importance(
        self,
        w_q: torch.Tensor | None,
        w_k: torch.Tensor | None,
        w_v: torch.Tensor | None,
    ) -> torch.Tensor:
        """
        Compute per-head importance as the sum of Frobenius norms of the
        corresponding Q, K, V slices.

        Returns: (num_heads,) tensor of importance scores.
        """
        H = self.num_heads
        D = self.head_dim
        device = None
        for w in (w_q, w_k, w_v):
            if w is not None:
                device = w.device
                break
        if device is None:
            device = torch.device("cpu")

        score = torch.zeros(H, device=device, dtype=torch.float32)

        if w_q is not None and w_q.shape[0] >= H * D:
            w_q_heads = w_q[: H * D].view(H, D, -1)
            score += w_q_heads.float().pow(2).sum(dim=(1, 2))

        if w_k is not None and w_k.shape[0] >= self.num_kv_heads * D:
            # K has num_kv_heads; expand each KV head to the Q heads it serves.
            w_k_heads = w_k[: self.num_kv_heads * D].view(self.num_kv_heads, D, -1)
            k_score = w_k_heads.float().pow(2).sum(dim=(1, 2))  # (num_kv_heads,)
            # Repeat each KV head `num_kv_groups` times to align with Q heads.
            rep = self.num_kv_groups
            k_expanded = k_score.repeat_interleave(rep)
            if k_expanded.numel() < H:
                k_expanded = k_expanded.repeat((H + k_expanded.numel() - 1) // k_expanded.numel())[
                    :H
                ]
            else:
                k_expanded = k_expanded[:H]
            score += k_expanded

        if w_v is not None and w_v.shape[0] >= self.num_kv_heads * D:
            w_v_heads = w_v[: self.num_kv_heads * D].view(self.num_kv_heads, D, -1)
            v_score = w_v_heads.float().pow(2).sum(dim=(1, 2))
            rep = self.num_kv_groups
            v_expanded = v_score.repeat_interleave(rep)
            if v_expanded.numel() < H:
                v_expanded = v_expanded.repeat((H + v_expanded.numel() - 1) // v_expanded.numel())[
                    :H
                ]
            else:
                v_expanded = v_expanded[:H]
            score += v_expanded

        return score

    # ------------------------------------------------------------------
    # copy_weights with head selection
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: nn.Module) -> None:
        with torch.no_grad():
            w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_out_lin = any(hasattr(old_attn, a) for a in ("out_lin", "c_proj", "o_proj"))
            if has_out_lin:
                self.is_nested_encoder = False
            elif (
                hasattr(old_attn, "self")
                and hasattr(old_attn.self, "query")
                and hasattr(old_attn, "output")
                and hasattr(old_attn.output, "dense")
            ) or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention"):
                self.is_nested_encoder = True
            else:
                self.is_nested_encoder = False

            # Determine nested container vs flat
            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    dense = old_attn.output.dense
                    w_o = dense.weight.detach()
                    b_o = dense.bias.detach() if dense.bias is not None else None
            else:
                src = old_attn

            # ---- Pattern 1: GPT-2 fused c_attn ----
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                p_cls = src.c_attn.__class__.__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(src.c_attn, "nf")
                    or (
                        c_attn_w.shape[0] == self.hidden_size
                        and c_attn_w.shape[1] != self.hidden_size
                    )
                ):
                    c_attn_w = c_attn_w.t().contiguous()

                H = self.hidden_size
                w_q = c_attn_w[:H, :]
                w_k = c_attn_w[H : 2 * H, :]
                w_v = c_attn_w[2 * H : 3 * H, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:H]
                    b_k = c_attn_b[H : 2 * H]
                    b_v = c_attn_b[2 * H : 3 * H]

                out_w = src.c_proj.weight.detach()
                p_proj_cls = src.c_proj.__class__.__name__
                if (
                    ("Conv1D" in p_proj_cls)
                    or hasattr(src.c_proj, "nf")
                    or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
                ):
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # ---- Pattern 2: BERT/RoBERTa nested ----
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # ---- Pattern 3: Qwen/LLaMA ----
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # ---- Pattern 4: DistilBERT ----
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # ---- Fallback: UniversalAdapter ----
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter(old_attn)
                    extracted = adapter.extract_qkv_out()
                    if extracted:
                        w_q = w_q if w_q is not None else extracted.get("q_weight")
                        b_q = b_q if b_q is not None else extracted.get("q_bias")
                        w_k = w_k if w_k is not None else extracted.get("k_weight")
                        b_k = b_k if b_k is not None else extracted.get("k_bias")
                        w_v = w_v if w_v is not None else extracted.get("v_weight")
                        b_v = b_v if b_v is not None else extracted.get("v_bias")
                        if w_o is None:
                            w_o = extracted.get("out_weight")
                            b_o = extracted.get("out_bias")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            missing = [n for n, w in (("q", w_q), ("k", w_k), ("v", w_v)) if w is None]
            if missing:
                raise RuntimeError(
                    f"SqueezingHeadsDistiller.copy_weights: missing mandatory "
                    f"projections {missing} from {src_cls_name}."
                )

            # Frobenius checks
            for name, w in (("w_q", w_q), ("w_k", w_k), ("w_v", w_v)):
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    logger.warning("SqueezingHeadsDistiller: %s near-zero Frobenius norm", name)

            # ---- Head importance & selection: SQ2-FIX-02 ----
            importance = self._compute_head_importance(w_q, w_k, w_v)  # (num_heads,)
            num_kept = self.num_kept_heads
            if num_kept < self.num_heads:
                kept_idx, _ = torch.topk(importance, num_kept, largest=True)
                kept_idx, _ = torch.sort(kept_idx)  # stable
            else:
                kept_idx = torch.arange(self.num_heads, device=importance.device)
            self.kept_head_indices = kept_idx.to(torch.long).cpu()

            logger.debug(
                "SqueezingHeadsDistiller: keeping %d/%d heads (indices=%s)",
                num_kept,
                self.num_heads,
                kept_idx.tolist(),
            )

            # ---- Build Q slices for kept heads ----
            D = self.head_dim
            H = self.num_heads
            w_q_heads = w_q[: H * D].view(H, D, -1)  # (H, D, hidden)
            b_q_heads = None
            if b_q is not None and b_q.shape[0] >= H * D:
                b_q_heads = b_q[: H * D].view(H, D)

            kept_q_w = w_q_heads[kept_idx].reshape(num_kept * D, w_q_heads.shape[-1])
            self.q_proj.weight.zero_()
            self.q_proj.weight.copy_(kept_q_w.to(self.q_proj.weight.dtype))

            if b_q_heads is not None and self.q_proj.bias is not None:
                kept_q_b = b_q_heads[kept_idx].reshape(num_kept * D)
                self.q_proj.bias.zero_()
                self.q_proj.bias.copy_(kept_q_b.to(self.q_proj.bias.dtype))

            # ---- K, V: copy ALL num_kv_heads (GQA-safe) ----
            self._copy_kv_full(self.k_proj, w_k, b_k)
            self._copy_kv_full(self.v_proj, w_v, b_v)

            # ---- Output projection ----
            # O must map from kept_q_head_size back to hidden.
            if self.is_nested_encoder:
                # Nested encoder bypasses out_proj entirely.
                nn.init.eye_(self.out_proj.weight)
                if self.out_proj.bias is not None:
                    nn.init.zeros_(self.out_proj.bias)
            elif w_o is not None:
                # w_o shape: (hidden, num_heads*D). Take kept head columns.
                if w_o.shape == (self.hidden_size, H * D):
                    w_o_heads = w_o.view(self.hidden_size, H, D)
                    kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                    self.out_proj.weight.zero_()
                    self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
                elif w_o.shape == (H * D, self.hidden_size):
                    # Transposed
                    w_o_t = w_o.t().contiguous()
                    w_o_heads = w_o_t.view(self.hidden_size, H, D)
                    kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                    self.out_proj.weight.zero_()
                    self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
                else:
                    raise RuntimeError(
                        f"SqueezingHeadsDistiller: unexpected w_o shape "
                        f"{tuple(w_o.shape)}; expected "
                        f"({self.hidden_size}, {H * D}) or ({H * D}, {self.hidden_size})."
                    )
                if b_o is not None and self.out_proj.bias is not None:
                    m = min(self.out_proj.bias.shape[0], b_o.shape[0])
                    self.out_proj.bias.zero_()
                    self.out_proj.bias[:m].copy_(b_o[:m].to(self.out_proj.bias.dtype))

    # ------------------------------------------------------------------
    # K/V copy helper: preserve ALL KV heads
    # ------------------------------------------------------------------
    def _copy_kv_full(self, target: LoRALinear, w: torch.Tensor, b: torch.Tensor | None) -> None:
        """Copy the full KV projection, sized to num_kv_heads * head_dim.

        Handles source having num_heads heads by mean-pooling to num_kv_heads
        (i.e., converting MHA → GQA). Never truncates heads silently.
        """
        target.weight.zero_()
        if b is not None and target.bias is not None:
            target.bias.zero_()

        target_out = target.weight.shape[0]
        target_in = target.weight.shape[1]
        D = self.head_dim

        # Case A: exact shape match
        if w.shape[0] == target_out and w.shape[1] == target_in:
            target.weight.copy_(w.to(target.weight.dtype))
            if b is not None and target.bias is not None and b.shape[0] == target_out:
                target.bias.copy_(b.to(target.bias.dtype))
            return

        # Case B: source has more heads (MHA → GQA pooling)
        src_out = w.shape[0]
        if src_out > target_out and src_out % D == 0 and target_out % D == 0:
            src_heads = src_out // D
            tgt_heads = target_out // D
            if src_heads % tgt_heads == 0:
                rep = src_heads // tgt_heads
                w_heads = w[:, :target_in].reshape(src_heads, D, target_in)
                w_pooled = w_heads.view(tgt_heads, rep, D, target_in).mean(dim=1)
                w_pooled = w_pooled.reshape(target_out, target_in)
                target.weight.copy_(w_pooled.to(target.weight.dtype))
                if b is not None and target.bias is not None:
                    b_heads = b.reshape(src_heads, D)
                    b_pooled = b_heads.view(tgt_heads, rep, D).mean(dim=1).reshape(target_out)
                    target.bias.copy_(b_pooled.to(target.bias.dtype))
                return

        # Case C: shape mismatch we cannot resolve
        raise RuntimeError(
            f"SqueezingHeadsDistiller: cannot map K/V weight from "
            f"{tuple(w.shape)} to target {tuple(target.weight.shape)}."
        )

    # ------------------------------------------------------------------
    # Mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(attention_mask, B, N_kv, device):
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # KV-cache helper
    # ------------------------------------------------------------------
    @staticmethod
    def _read_past_kv(past_key_value, layer_idx=0):
        if past_key_value is None:
            return None, None, 0
        if hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
            idx = min(layer_idx, len(past_key_value.key_cache) - 1)
            k = past_key_value.key_cache[idx]
            v = past_key_value.value_cache[idx]
            return k, v, k.shape[-2]
        if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
            k, v = past_key_value[0], past_key_value[1]
            return k, v, k.shape[-2]
        return None, None, 0

    # ------------------------------------------------------------------
    # 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, L, _ = 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
        L_kv_new = kv_input.shape[1]

        # ---- Projections ----
        q = (
            self.q_proj(hidden_states)
            .view(B, L, self.num_kept_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k_new = (
            self.k_proj(kv_input)
            .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v_new = (
            self.v_proj(kv_input)
            .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- RoPE ----
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        past_k_probe, _, past_len_rope = self._read_past_kv(past_key_value, layer_idx)
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            cos = cos[..., past_len_rope : past_len_rope + L, :]
            sin = sin[..., past_len_rope : past_len_rope + L, :]
            q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin)

        # ---- KV cache ----
        past_k, past_v, past_len = self._read_past_kv(past_key_value, layer_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]

        # ---- GQA expansion ----
        # kept Q heads map to KV heads via (kept_idx // num_kv_groups).
        # Here we simply expand KV to num_kv_groups * num_kv_heads, then select
        # the subset corresponding to kept heads.
        if self.num_kv_heads < self.num_heads:
            # Expand KV to full num_heads then gather kept heads.
            k_full = repeat_kv(k_cat, self.num_kv_groups)  # (B, num_heads, N_kv, D)
            v_full = repeat_kv(v_cat, self.num_kv_groups)
            if self.kept_head_indices is not None and self.num_kept_heads < self.num_heads:
                idx = self.kept_head_indices.to(k_full.device)
                k_full = k_full.index_select(1, idx)
                v_full = v_full.index_select(1, idx)
        else:
            k_full = k_cat
            v_full = v_cat

        # ---- Mask ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

        scale = 1.0 / math.sqrt(self.head_dim)

        # ---- SDPA fast path ----
        use_sdpa = not output_attentions and head_mask is None
        if use_sdpa:
            q_sdp = q.to(torch.float16) if q.dtype == torch.float32 else q
            k_sdp = k_full.to(torch.float16) if k_full.dtype == torch.float32 else k_full
            v_sdp = v_full.to(torch.float16) if v_full.dtype == torch.float32 else v_full

            attn_bias_sdp = None
            if token_valid is not None:
                attn_bias_sdp = ((token_valid[:, None, None, :].float() - 1.0) * 1e4).to(
                    q_sdp.dtype
                )

            is_causal_sdp = self.is_decoder and not is_cross and past_len == 0 and L > 1

            try:
                attn_out = 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,
                )
                attn_out = (
                    attn_out.to(hidden_states.dtype)
                    .transpose(1, 2)
                    .contiguous()
                    .view(B, L, self.kept_q_head_size)
                )
                attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
                attn_probs_ret = None
            except Exception as e:
                logger.debug("SDPA failed (%s); manual path.", e)
                use_sdpa = False

        if not use_sdpa:
            scores = torch.matmul(q.float(), k_full.float().transpose(-2, -1).contiguous()) * scale

            if self.is_decoder and not is_cross:
                if past_len > 0:
                    abs_q = (past_len + torch.arange(L, device=device)).unsqueeze(1)
                    abs_k = torch.arange(N_kv, device=device).unsqueeze(0)
                    causal = abs_k <= abs_q
                else:
                    causal = torch.tril(torch.ones(L, N_kv, device=device, dtype=torch.bool))
                scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

            if token_valid is not None:
                scores = scores.masked_fill(~token_valid[:, None, None, :].bool(), -1e4)

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(hidden_states.dtype)
            probs = self.dropout(probs)

            # head_mask: apply to (B, num_kept_heads, 1, 1)
            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
                if hm.shape[1] != self.num_kept_heads and hm.shape[1] == self.num_heads:
                    # head_mask is for the FULL head count; select kept heads.
                    if self.kept_head_indices is not None:
                        idx = self.kept_head_indices.to(hm.device)
                        hm = hm.index_select(1, idx)
                probs = probs * hm.to(probs.dtype)

            attn_out = (
                torch.matmul(probs, v_full)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.kept_q_head_size)
            )
            attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
            attn_probs_ret = probs if output_attentions else None
        else:
            attn_probs_ret = None

        # ---- Output ----
        if self.is_nested_encoder:
            out = attn_out
        else:
            out = self.out_proj(attn_out)

        if out.shape[-1] != hidden_states.shape[-1]:
            raise RuntimeError(
                f"SqueezingHeadsDistiller: output dim {out.shape[-1]} != "
                f"input dim {hidden_states.shape[-1]}."
            )

        return out, attn_probs_ret

Attention wrapper that keeps only the top-k most important heads (based on pretrained Q/K/V Frobenius norms) while preserving GQA structure.

Effective head count: num_kept_heads = round(head_keep_ratio * num_heads) q_head_size = num_kept_heads * head_dim out_head_size = num_kept_heads * head_dim

KV projections preserve all num_kv_heads (GQA). Kept Q heads are grouped to matching KV heads via GQA index mapping.

Parameter reduction: Q : num_heads * head_dim * hidden → num_kept_heads * head_dim * hidden O : hidden * num_heads * head_dim → hidden * num_kept_heads * head_dim K/V: unchanged (GQA-compatible)

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.

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: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    with torch.no_grad():
        w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        has_out_lin = any(hasattr(old_attn, a) for a in ("out_lin", "c_proj", "o_proj"))
        if has_out_lin:
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention"):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        # Determine nested container vs flat
        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                dense = old_attn.output.dense
                w_o = dense.weight.detach()
                b_o = dense.bias.detach() if dense.bias is not None else None
        else:
            src = old_attn

        # ---- Pattern 1: GPT-2 fused c_attn ----
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
            p_cls = src.c_attn.__class__.__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(src.c_attn, "nf")
                or (
                    c_attn_w.shape[0] == self.hidden_size
                    and c_attn_w.shape[1] != self.hidden_size
                )
            ):
                c_attn_w = c_attn_w.t().contiguous()

            H = self.hidden_size
            w_q = c_attn_w[:H, :]
            w_k = c_attn_w[H : 2 * H, :]
            w_v = c_attn_w[2 * H : 3 * H, :]
            if c_attn_b is not None:
                b_q = c_attn_b[:H]
                b_k = c_attn_b[H : 2 * H]
                b_v = c_attn_b[2 * H : 3 * H]

            out_w = src.c_proj.weight.detach()
            p_proj_cls = src.c_proj.__class__.__name__
            if (
                ("Conv1D" in p_proj_cls)
                or hasattr(src.c_proj, "nf")
                or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
            ):
                out_w = out_w.t().contiguous()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # ---- Pattern 2: BERT/RoBERTa nested ----
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None
            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # ---- Pattern 3: Qwen/LLaMA ----
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # ---- Pattern 4: DistilBERT ----
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # ---- Fallback: UniversalAdapter ----
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter(old_attn)
                extracted = adapter.extract_qkv_out()
                if extracted:
                    w_q = w_q if w_q is not None else extracted.get("q_weight")
                    b_q = b_q if b_q is not None else extracted.get("q_bias")
                    w_k = w_k if w_k is not None else extracted.get("k_weight")
                    b_k = b_k if b_k is not None else extracted.get("k_bias")
                    w_v = w_v if w_v is not None else extracted.get("v_weight")
                    b_v = b_v if b_v is not None else extracted.get("v_bias")
                    if w_o is None:
                        w_o = extracted.get("out_weight")
                        b_o = extracted.get("out_bias")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        missing = [n for n, w in (("q", w_q), ("k", w_k), ("v", w_v)) if w is None]
        if missing:
            raise RuntimeError(
                f"SqueezingHeadsDistiller.copy_weights: missing mandatory "
                f"projections {missing} from {src_cls_name}."
            )

        # Frobenius checks
        for name, w in (("w_q", w_q), ("w_k", w_k), ("w_v", w_v)):
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                logger.warning("SqueezingHeadsDistiller: %s near-zero Frobenius norm", name)

        # ---- Head importance & selection: SQ2-FIX-02 ----
        importance = self._compute_head_importance(w_q, w_k, w_v)  # (num_heads,)
        num_kept = self.num_kept_heads
        if num_kept < self.num_heads:
            kept_idx, _ = torch.topk(importance, num_kept, largest=True)
            kept_idx, _ = torch.sort(kept_idx)  # stable
        else:
            kept_idx = torch.arange(self.num_heads, device=importance.device)
        self.kept_head_indices = kept_idx.to(torch.long).cpu()

        logger.debug(
            "SqueezingHeadsDistiller: keeping %d/%d heads (indices=%s)",
            num_kept,
            self.num_heads,
            kept_idx.tolist(),
        )

        # ---- Build Q slices for kept heads ----
        D = self.head_dim
        H = self.num_heads
        w_q_heads = w_q[: H * D].view(H, D, -1)  # (H, D, hidden)
        b_q_heads = None
        if b_q is not None and b_q.shape[0] >= H * D:
            b_q_heads = b_q[: H * D].view(H, D)

        kept_q_w = w_q_heads[kept_idx].reshape(num_kept * D, w_q_heads.shape[-1])
        self.q_proj.weight.zero_()
        self.q_proj.weight.copy_(kept_q_w.to(self.q_proj.weight.dtype))

        if b_q_heads is not None and self.q_proj.bias is not None:
            kept_q_b = b_q_heads[kept_idx].reshape(num_kept * D)
            self.q_proj.bias.zero_()
            self.q_proj.bias.copy_(kept_q_b.to(self.q_proj.bias.dtype))

        # ---- K, V: copy ALL num_kv_heads (GQA-safe) ----
        self._copy_kv_full(self.k_proj, w_k, b_k)
        self._copy_kv_full(self.v_proj, w_v, b_v)

        # ---- Output projection ----
        # O must map from kept_q_head_size back to hidden.
        if self.is_nested_encoder:
            # Nested encoder bypasses out_proj entirely.
            nn.init.eye_(self.out_proj.weight)
            if self.out_proj.bias is not None:
                nn.init.zeros_(self.out_proj.bias)
        elif w_o is not None:
            # w_o shape: (hidden, num_heads*D). Take kept head columns.
            if w_o.shape == (self.hidden_size, H * D):
                w_o_heads = w_o.view(self.hidden_size, H, D)
                kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                self.out_proj.weight.zero_()
                self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
            elif w_o.shape == (H * D, self.hidden_size):
                # Transposed
                w_o_t = w_o.t().contiguous()
                w_o_heads = w_o_t.view(self.hidden_size, H, D)
                kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                self.out_proj.weight.zero_()
                self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
            else:
                raise RuntimeError(
                    f"SqueezingHeadsDistiller: unexpected w_o shape "
                    f"{tuple(w_o.shape)}; expected "
                    f"({self.hidden_size}, {H * D}) or ({H * D}, {self.hidden_size})."
                )
            if b_o is not None and self.out_proj.bias is not None:
                m = min(self.out_proj.bias.shape[0], b_o.shape[0])
                self.out_proj.bias.zero_()
                self.out_proj.bias[:m].copy_(b_o[:m].to(self.out_proj.bias.dtype))

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, L, _ = 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
    L_kv_new = kv_input.shape[1]

    # ---- Projections ----
    q = (
        self.q_proj(hidden_states)
        .view(B, L, self.num_kept_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    k_new = (
        self.k_proj(kv_input)
        .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v_new = (
        self.v_proj(kv_input)
        .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # ---- RoPE ----
    if position_embeddings is None:
        position_embeddings = kwargs.get("position_embeddings")
    past_k_probe, _, past_len_rope = self._read_past_kv(past_key_value, layer_idx)
    if position_embeddings is not None and not is_cross:
        cos, sin = position_embeddings
        cos = cos[..., past_len_rope : past_len_rope + L, :]
        sin = sin[..., past_len_rope : past_len_rope + L, :]
        q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin)

    # ---- KV cache ----
    past_k, past_v, past_len = self._read_past_kv(past_key_value, layer_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]

    # ---- GQA expansion ----
    # kept Q heads map to KV heads via (kept_idx // num_kv_groups).
    # Here we simply expand KV to num_kv_groups * num_kv_heads, then select
    # the subset corresponding to kept heads.
    if self.num_kv_heads < self.num_heads:
        # Expand KV to full num_heads then gather kept heads.
        k_full = repeat_kv(k_cat, self.num_kv_groups)  # (B, num_heads, N_kv, D)
        v_full = repeat_kv(v_cat, self.num_kv_groups)
        if self.kept_head_indices is not None and self.num_kept_heads < self.num_heads:
            idx = self.kept_head_indices.to(k_full.device)
            k_full = k_full.index_select(1, idx)
            v_full = v_full.index_select(1, idx)
    else:
        k_full = k_cat
        v_full = v_cat

    # ---- Mask ----
    eff_mask = encoder_attention_mask if is_cross else attention_mask
    token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

    scale = 1.0 / math.sqrt(self.head_dim)

    # ---- SDPA fast path ----
    use_sdpa = not output_attentions and head_mask is None
    if use_sdpa:
        q_sdp = q.to(torch.float16) if q.dtype == torch.float32 else q
        k_sdp = k_full.to(torch.float16) if k_full.dtype == torch.float32 else k_full
        v_sdp = v_full.to(torch.float16) if v_full.dtype == torch.float32 else v_full

        attn_bias_sdp = None
        if token_valid is not None:
            attn_bias_sdp = ((token_valid[:, None, None, :].float() - 1.0) * 1e4).to(
                q_sdp.dtype
            )

        is_causal_sdp = self.is_decoder and not is_cross and past_len == 0 and L > 1

        try:
            attn_out = 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,
            )
            attn_out = (
                attn_out.to(hidden_states.dtype)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.kept_q_head_size)
            )
            attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
            attn_probs_ret = None
        except Exception as e:
            logger.debug("SDPA failed (%s); manual path.", e)
            use_sdpa = False

    if not use_sdpa:
        scores = torch.matmul(q.float(), k_full.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross:
            if past_len > 0:
                abs_q = (past_len + torch.arange(L, device=device)).unsqueeze(1)
                abs_k = torch.arange(N_kv, device=device).unsqueeze(0)
                causal = abs_k <= abs_q
            else:
                causal = torch.tril(torch.ones(L, N_kv, device=device, dtype=torch.bool))
            scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

        if token_valid is not None:
            scores = scores.masked_fill(~token_valid[:, None, None, :].bool(), -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        probs = self.dropout(probs)

        # head_mask: apply to (B, num_kept_heads, 1, 1)
        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
            if hm.shape[1] != self.num_kept_heads and hm.shape[1] == self.num_heads:
                # head_mask is for the FULL head count; select kept heads.
                if self.kept_head_indices is not None:
                    idx = self.kept_head_indices.to(hm.device)
                    hm = hm.index_select(1, idx)
            probs = probs * hm.to(probs.dtype)

        attn_out = (
            torch.matmul(probs, v_full)
            .transpose(1, 2)
            .contiguous()
            .view(B, L, self.kept_q_head_size)
        )
        attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
        attn_probs_ret = probs if output_attentions else None
    else:
        attn_probs_ret = None

    # ---- Output ----
    if self.is_nested_encoder:
        out = attn_out
    else:
        out = self.out_proj(attn_out)

    if out.shape[-1] != hidden_states.shape[-1]:
        raise RuntimeError(
            f"SqueezingHeadsDistiller: output dim {out.shape[-1]} != "
            f"input dim {hidden_states.shape[-1]}."
        )

    return out, attn_probs_ret

Must return (output, attention_weights_or_None).

class SqueezingHeadsAttention (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
head_keep_ratio: float | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("squeezing_heads_distiller", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("squeezing_heads_distill", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("squeezing_heads", supports_causal=True, supports_gqa=True)
class SqueezingHeadsDistiller(BaseAttention):
    """
    Attention wrapper that keeps only the top-k most important heads
    (based on pretrained Q/K/V Frobenius norms) while preserving GQA structure.

    Effective head count:
        num_kept_heads = round(head_keep_ratio * num_heads)
        q_head_size   = num_kept_heads * head_dim
        out_head_size = num_kept_heads * head_dim

    KV projections preserve all num_kv_heads (GQA). Kept Q heads are grouped
    to matching KV heads via GQA index mapping.

    Parameter reduction:
        Q : num_heads   * head_dim * hidden  →  num_kept_heads * head_dim * hidden
        O : hidden * num_heads * head_dim    →  hidden * num_kept_heads * head_dim
        K/V: unchanged (GQA-compatible)
    """

    supports_gqa: bool = True
    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = None,
        head_keep_ratio: float | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        # ---- 1. Dimensions ----
        self.hidden_size = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.num_heads = int(
            _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "num_heads", kwargs.get("num_attention_heads", 12)),
            )
        )

        cfg_hd = _cfg_get(config, "head_dim", None)
        kw_hd = kwargs.get("head_dim")
        if cfg_hd is not None:
            self.head_dim = int(cfg_hd)
        elif kw_hd is not None:
            self.head_dim = int(kw_hd)
        elif old_attn is not None:
            self.head_dim = None
            for attr in ("q_proj", "query", "q_lin"):
                if hasattr(old_attn, attr):
                    qw = getattr(old_attn, attr).weight
                    self.head_dim = qw.shape[0] // max(1, self.num_heads)
                    break
            if self.head_dim is None:
                self.head_dim = self.hidden_size // max(1, self.num_heads)
        else:
            self.head_dim = self.hidden_size // max(1, self.num_heads)

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))

        # ---- 2. Compute num_kept_heads ----
        # SQ2-FIX-01: derive head_keep_ratio from `rank` if not explicitly given.
        if head_keep_ratio is not None:
            hkr = float(head_keep_ratio)
        elif rank is not None and rank > 0:
            # If rank <= num_heads, interpret as "keep this many heads"
            if rank <= self.num_heads:
                hkr = float(rank) / float(self.num_heads)
            else:
                # Otherwise interpret as "target head_dim"
                effective_hd = min(int(rank), self.head_dim)
                hkr = float(effective_hd) / float(self.head_dim)
        else:
            hkr = 1.0
        self.head_keep_ratio = max(1.0 / self.num_heads, min(1.0, hkr))

        self.num_kept_heads = max(
            1, min(self.num_heads, int(round(self.num_heads * self.head_keep_ratio)))
        )
        self.head_keep_ratio = self.num_kept_heads / self.num_heads

        self.kept_q_head_size = self.num_kept_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim
        self.all_head_size = self.kept_q_head_size

        # GQA mapping: which KV head does each kept Q head map to?
        # Kept heads are the top-k by importance; we must remember their original indices.
        self.num_kv_groups = max(1, self.num_heads // self.num_kv_heads)
        self.kept_head_indices: torch.Tensor | None = None  # filled by copy_weights

        # ---- 3. Decoder detection ----
        self.is_decoder = _detect_decoder(config, is_decoder, old_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        self.lora_rank = int(lora_rank if lora_rank is not None else 16)
        self.lora_alpha = float(lora_alpha)

        # ---- 4. Projections ----
        # SQ2-FIX-01: Q and O are sized to num_kept_heads, not num_heads.
        self.q_proj = LoRALinear(
            self.hidden_size,
            self.kept_q_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.kept_q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        # ---- 5. Copy weights from source ----
        src_attn = old_attn if old_attn is not None else kwargs.get("pretrained_attn")
        if src_attn is not None:
            self.copy_weights(src_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 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

    # ------------------------------------------------------------------
    # Head importance: SQ2-FIX-02
    # ------------------------------------------------------------------
    def _compute_head_importance(
        self,
        w_q: torch.Tensor | None,
        w_k: torch.Tensor | None,
        w_v: torch.Tensor | None,
    ) -> torch.Tensor:
        """
        Compute per-head importance as the sum of Frobenius norms of the
        corresponding Q, K, V slices.

        Returns: (num_heads,) tensor of importance scores.
        """
        H = self.num_heads
        D = self.head_dim
        device = None
        for w in (w_q, w_k, w_v):
            if w is not None:
                device = w.device
                break
        if device is None:
            device = torch.device("cpu")

        score = torch.zeros(H, device=device, dtype=torch.float32)

        if w_q is not None and w_q.shape[0] >= H * D:
            w_q_heads = w_q[: H * D].view(H, D, -1)
            score += w_q_heads.float().pow(2).sum(dim=(1, 2))

        if w_k is not None and w_k.shape[0] >= self.num_kv_heads * D:
            # K has num_kv_heads; expand each KV head to the Q heads it serves.
            w_k_heads = w_k[: self.num_kv_heads * D].view(self.num_kv_heads, D, -1)
            k_score = w_k_heads.float().pow(2).sum(dim=(1, 2))  # (num_kv_heads,)
            # Repeat each KV head `num_kv_groups` times to align with Q heads.
            rep = self.num_kv_groups
            k_expanded = k_score.repeat_interleave(rep)
            if k_expanded.numel() < H:
                k_expanded = k_expanded.repeat((H + k_expanded.numel() - 1) // k_expanded.numel())[
                    :H
                ]
            else:
                k_expanded = k_expanded[:H]
            score += k_expanded

        if w_v is not None and w_v.shape[0] >= self.num_kv_heads * D:
            w_v_heads = w_v[: self.num_kv_heads * D].view(self.num_kv_heads, D, -1)
            v_score = w_v_heads.float().pow(2).sum(dim=(1, 2))
            rep = self.num_kv_groups
            v_expanded = v_score.repeat_interleave(rep)
            if v_expanded.numel() < H:
                v_expanded = v_expanded.repeat((H + v_expanded.numel() - 1) // v_expanded.numel())[
                    :H
                ]
            else:
                v_expanded = v_expanded[:H]
            score += v_expanded

        return score

    # ------------------------------------------------------------------
    # copy_weights with head selection
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: nn.Module) -> None:
        with torch.no_grad():
            w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_out_lin = any(hasattr(old_attn, a) for a in ("out_lin", "c_proj", "o_proj"))
            if has_out_lin:
                self.is_nested_encoder = False
            elif (
                hasattr(old_attn, "self")
                and hasattr(old_attn.self, "query")
                and hasattr(old_attn, "output")
                and hasattr(old_attn.output, "dense")
            ) or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention"):
                self.is_nested_encoder = True
            else:
                self.is_nested_encoder = False

            # Determine nested container vs flat
            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    dense = old_attn.output.dense
                    w_o = dense.weight.detach()
                    b_o = dense.bias.detach() if dense.bias is not None else None
            else:
                src = old_attn

            # ---- Pattern 1: GPT-2 fused c_attn ----
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                p_cls = src.c_attn.__class__.__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(src.c_attn, "nf")
                    or (
                        c_attn_w.shape[0] == self.hidden_size
                        and c_attn_w.shape[1] != self.hidden_size
                    )
                ):
                    c_attn_w = c_attn_w.t().contiguous()

                H = self.hidden_size
                w_q = c_attn_w[:H, :]
                w_k = c_attn_w[H : 2 * H, :]
                w_v = c_attn_w[2 * H : 3 * H, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:H]
                    b_k = c_attn_b[H : 2 * H]
                    b_v = c_attn_b[2 * H : 3 * H]

                out_w = src.c_proj.weight.detach()
                p_proj_cls = src.c_proj.__class__.__name__
                if (
                    ("Conv1D" in p_proj_cls)
                    or hasattr(src.c_proj, "nf")
                    or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
                ):
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # ---- Pattern 2: BERT/RoBERTa nested ----
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # ---- Pattern 3: Qwen/LLaMA ----
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # ---- Pattern 4: DistilBERT ----
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # ---- Fallback: UniversalAdapter ----
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter(old_attn)
                    extracted = adapter.extract_qkv_out()
                    if extracted:
                        w_q = w_q if w_q is not None else extracted.get("q_weight")
                        b_q = b_q if b_q is not None else extracted.get("q_bias")
                        w_k = w_k if w_k is not None else extracted.get("k_weight")
                        b_k = b_k if b_k is not None else extracted.get("k_bias")
                        w_v = w_v if w_v is not None else extracted.get("v_weight")
                        b_v = b_v if b_v is not None else extracted.get("v_bias")
                        if w_o is None:
                            w_o = extracted.get("out_weight")
                            b_o = extracted.get("out_bias")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            missing = [n for n, w in (("q", w_q), ("k", w_k), ("v", w_v)) if w is None]
            if missing:
                raise RuntimeError(
                    f"SqueezingHeadsDistiller.copy_weights: missing mandatory "
                    f"projections {missing} from {src_cls_name}."
                )

            # Frobenius checks
            for name, w in (("w_q", w_q), ("w_k", w_k), ("w_v", w_v)):
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    logger.warning("SqueezingHeadsDistiller: %s near-zero Frobenius norm", name)

            # ---- Head importance & selection: SQ2-FIX-02 ----
            importance = self._compute_head_importance(w_q, w_k, w_v)  # (num_heads,)
            num_kept = self.num_kept_heads
            if num_kept < self.num_heads:
                kept_idx, _ = torch.topk(importance, num_kept, largest=True)
                kept_idx, _ = torch.sort(kept_idx)  # stable
            else:
                kept_idx = torch.arange(self.num_heads, device=importance.device)
            self.kept_head_indices = kept_idx.to(torch.long).cpu()

            logger.debug(
                "SqueezingHeadsDistiller: keeping %d/%d heads (indices=%s)",
                num_kept,
                self.num_heads,
                kept_idx.tolist(),
            )

            # ---- Build Q slices for kept heads ----
            D = self.head_dim
            H = self.num_heads
            w_q_heads = w_q[: H * D].view(H, D, -1)  # (H, D, hidden)
            b_q_heads = None
            if b_q is not None and b_q.shape[0] >= H * D:
                b_q_heads = b_q[: H * D].view(H, D)

            kept_q_w = w_q_heads[kept_idx].reshape(num_kept * D, w_q_heads.shape[-1])
            self.q_proj.weight.zero_()
            self.q_proj.weight.copy_(kept_q_w.to(self.q_proj.weight.dtype))

            if b_q_heads is not None and self.q_proj.bias is not None:
                kept_q_b = b_q_heads[kept_idx].reshape(num_kept * D)
                self.q_proj.bias.zero_()
                self.q_proj.bias.copy_(kept_q_b.to(self.q_proj.bias.dtype))

            # ---- K, V: copy ALL num_kv_heads (GQA-safe) ----
            self._copy_kv_full(self.k_proj, w_k, b_k)
            self._copy_kv_full(self.v_proj, w_v, b_v)

            # ---- Output projection ----
            # O must map from kept_q_head_size back to hidden.
            if self.is_nested_encoder:
                # Nested encoder bypasses out_proj entirely.
                nn.init.eye_(self.out_proj.weight)
                if self.out_proj.bias is not None:
                    nn.init.zeros_(self.out_proj.bias)
            elif w_o is not None:
                # w_o shape: (hidden, num_heads*D). Take kept head columns.
                if w_o.shape == (self.hidden_size, H * D):
                    w_o_heads = w_o.view(self.hidden_size, H, D)
                    kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                    self.out_proj.weight.zero_()
                    self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
                elif w_o.shape == (H * D, self.hidden_size):
                    # Transposed
                    w_o_t = w_o.t().contiguous()
                    w_o_heads = w_o_t.view(self.hidden_size, H, D)
                    kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                    self.out_proj.weight.zero_()
                    self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
                else:
                    raise RuntimeError(
                        f"SqueezingHeadsDistiller: unexpected w_o shape "
                        f"{tuple(w_o.shape)}; expected "
                        f"({self.hidden_size}, {H * D}) or ({H * D}, {self.hidden_size})."
                    )
                if b_o is not None and self.out_proj.bias is not None:
                    m = min(self.out_proj.bias.shape[0], b_o.shape[0])
                    self.out_proj.bias.zero_()
                    self.out_proj.bias[:m].copy_(b_o[:m].to(self.out_proj.bias.dtype))

    # ------------------------------------------------------------------
    # K/V copy helper: preserve ALL KV heads
    # ------------------------------------------------------------------
    def _copy_kv_full(self, target: LoRALinear, w: torch.Tensor, b: torch.Tensor | None) -> None:
        """Copy the full KV projection, sized to num_kv_heads * head_dim.

        Handles source having num_heads heads by mean-pooling to num_kv_heads
        (i.e., converting MHA → GQA). Never truncates heads silently.
        """
        target.weight.zero_()
        if b is not None and target.bias is not None:
            target.bias.zero_()

        target_out = target.weight.shape[0]
        target_in = target.weight.shape[1]
        D = self.head_dim

        # Case A: exact shape match
        if w.shape[0] == target_out and w.shape[1] == target_in:
            target.weight.copy_(w.to(target.weight.dtype))
            if b is not None and target.bias is not None and b.shape[0] == target_out:
                target.bias.copy_(b.to(target.bias.dtype))
            return

        # Case B: source has more heads (MHA → GQA pooling)
        src_out = w.shape[0]
        if src_out > target_out and src_out % D == 0 and target_out % D == 0:
            src_heads = src_out // D
            tgt_heads = target_out // D
            if src_heads % tgt_heads == 0:
                rep = src_heads // tgt_heads
                w_heads = w[:, :target_in].reshape(src_heads, D, target_in)
                w_pooled = w_heads.view(tgt_heads, rep, D, target_in).mean(dim=1)
                w_pooled = w_pooled.reshape(target_out, target_in)
                target.weight.copy_(w_pooled.to(target.weight.dtype))
                if b is not None and target.bias is not None:
                    b_heads = b.reshape(src_heads, D)
                    b_pooled = b_heads.view(tgt_heads, rep, D).mean(dim=1).reshape(target_out)
                    target.bias.copy_(b_pooled.to(target.bias.dtype))
                return

        # Case C: shape mismatch we cannot resolve
        raise RuntimeError(
            f"SqueezingHeadsDistiller: cannot map K/V weight from "
            f"{tuple(w.shape)} to target {tuple(target.weight.shape)}."
        )

    # ------------------------------------------------------------------
    # Mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(attention_mask, B, N_kv, device):
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # KV-cache helper
    # ------------------------------------------------------------------
    @staticmethod
    def _read_past_kv(past_key_value, layer_idx=0):
        if past_key_value is None:
            return None, None, 0
        if hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
            idx = min(layer_idx, len(past_key_value.key_cache) - 1)
            k = past_key_value.key_cache[idx]
            v = past_key_value.value_cache[idx]
            return k, v, k.shape[-2]
        if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
            k, v = past_key_value[0], past_key_value[1]
            return k, v, k.shape[-2]
        return None, None, 0

    # ------------------------------------------------------------------
    # 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, L, _ = 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
        L_kv_new = kv_input.shape[1]

        # ---- Projections ----
        q = (
            self.q_proj(hidden_states)
            .view(B, L, self.num_kept_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k_new = (
            self.k_proj(kv_input)
            .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v_new = (
            self.v_proj(kv_input)
            .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- RoPE ----
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        past_k_probe, _, past_len_rope = self._read_past_kv(past_key_value, layer_idx)
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            cos = cos[..., past_len_rope : past_len_rope + L, :]
            sin = sin[..., past_len_rope : past_len_rope + L, :]
            q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin)

        # ---- KV cache ----
        past_k, past_v, past_len = self._read_past_kv(past_key_value, layer_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]

        # ---- GQA expansion ----
        # kept Q heads map to KV heads via (kept_idx // num_kv_groups).
        # Here we simply expand KV to num_kv_groups * num_kv_heads, then select
        # the subset corresponding to kept heads.
        if self.num_kv_heads < self.num_heads:
            # Expand KV to full num_heads then gather kept heads.
            k_full = repeat_kv(k_cat, self.num_kv_groups)  # (B, num_heads, N_kv, D)
            v_full = repeat_kv(v_cat, self.num_kv_groups)
            if self.kept_head_indices is not None and self.num_kept_heads < self.num_heads:
                idx = self.kept_head_indices.to(k_full.device)
                k_full = k_full.index_select(1, idx)
                v_full = v_full.index_select(1, idx)
        else:
            k_full = k_cat
            v_full = v_cat

        # ---- Mask ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

        scale = 1.0 / math.sqrt(self.head_dim)

        # ---- SDPA fast path ----
        use_sdpa = not output_attentions and head_mask is None
        if use_sdpa:
            q_sdp = q.to(torch.float16) if q.dtype == torch.float32 else q
            k_sdp = k_full.to(torch.float16) if k_full.dtype == torch.float32 else k_full
            v_sdp = v_full.to(torch.float16) if v_full.dtype == torch.float32 else v_full

            attn_bias_sdp = None
            if token_valid is not None:
                attn_bias_sdp = ((token_valid[:, None, None, :].float() - 1.0) * 1e4).to(
                    q_sdp.dtype
                )

            is_causal_sdp = self.is_decoder and not is_cross and past_len == 0 and L > 1

            try:
                attn_out = 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,
                )
                attn_out = (
                    attn_out.to(hidden_states.dtype)
                    .transpose(1, 2)
                    .contiguous()
                    .view(B, L, self.kept_q_head_size)
                )
                attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
                attn_probs_ret = None
            except Exception as e:
                logger.debug("SDPA failed (%s); manual path.", e)
                use_sdpa = False

        if not use_sdpa:
            scores = torch.matmul(q.float(), k_full.float().transpose(-2, -1).contiguous()) * scale

            if self.is_decoder and not is_cross:
                if past_len > 0:
                    abs_q = (past_len + torch.arange(L, device=device)).unsqueeze(1)
                    abs_k = torch.arange(N_kv, device=device).unsqueeze(0)
                    causal = abs_k <= abs_q
                else:
                    causal = torch.tril(torch.ones(L, N_kv, device=device, dtype=torch.bool))
                scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

            if token_valid is not None:
                scores = scores.masked_fill(~token_valid[:, None, None, :].bool(), -1e4)

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(hidden_states.dtype)
            probs = self.dropout(probs)

            # head_mask: apply to (B, num_kept_heads, 1, 1)
            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
                if hm.shape[1] != self.num_kept_heads and hm.shape[1] == self.num_heads:
                    # head_mask is for the FULL head count; select kept heads.
                    if self.kept_head_indices is not None:
                        idx = self.kept_head_indices.to(hm.device)
                        hm = hm.index_select(1, idx)
                probs = probs * hm.to(probs.dtype)

            attn_out = (
                torch.matmul(probs, v_full)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.kept_q_head_size)
            )
            attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
            attn_probs_ret = probs if output_attentions else None
        else:
            attn_probs_ret = None

        # ---- Output ----
        if self.is_nested_encoder:
            out = attn_out
        else:
            out = self.out_proj(attn_out)

        if out.shape[-1] != hidden_states.shape[-1]:
            raise RuntimeError(
                f"SqueezingHeadsDistiller: output dim {out.shape[-1]} != "
                f"input dim {hidden_states.shape[-1]}."
            )

        return out, attn_probs_ret

Attention wrapper that keeps only the top-k most important heads (based on pretrained Q/K/V Frobenius norms) while preserving GQA structure.

Effective head count: num_kept_heads = round(head_keep_ratio * num_heads) q_head_size = num_kept_heads * head_dim out_head_size = num_kept_heads * head_dim

KV projections preserve all num_kv_heads (GQA). Kept Q heads are grouped to matching KV heads via GQA index mapping.

Parameter reduction: Q : num_heads * head_dim * hidden → num_kept_heads * head_dim * hidden O : hidden * num_heads * head_dim → hidden * num_kept_heads * head_dim K/V: unchanged (GQA-compatible)

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.

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: torch.nn.modules.module.Module) ‑> None

Expand source code

def copy_weights(self, old_attn: nn.Module) -> None:
    with torch.no_grad():
        w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        has_out_lin = any(hasattr(old_attn, a) for a in ("out_lin", "c_proj", "o_proj"))
        if has_out_lin:
            self.is_nested_encoder = False
        elif (
            hasattr(old_attn, "self")
            and hasattr(old_attn.self, "query")
            and hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
        ) or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention"):
            self.is_nested_encoder = True
        else:
            self.is_nested_encoder = False

        # Determine nested container vs flat
        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                dense = old_attn.output.dense
                w_o = dense.weight.detach()
                b_o = dense.bias.detach() if dense.bias is not None else None
        else:
            src = old_attn

        # ---- Pattern 1: GPT-2 fused c_attn ----
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
            p_cls = src.c_attn.__class__.__name__
            if (
                ("Conv1D" in p_cls)
                or hasattr(src.c_attn, "nf")
                or (
                    c_attn_w.shape[0] == self.hidden_size
                    and c_attn_w.shape[1] != self.hidden_size
                )
            ):
                c_attn_w = c_attn_w.t().contiguous()

            H = self.hidden_size
            w_q = c_attn_w[:H, :]
            w_k = c_attn_w[H : 2 * H, :]
            w_v = c_attn_w[2 * H : 3 * H, :]
            if c_attn_b is not None:
                b_q = c_attn_b[:H]
                b_k = c_attn_b[H : 2 * H]
                b_v = c_attn_b[2 * H : 3 * H]

            out_w = src.c_proj.weight.detach()
            p_proj_cls = src.c_proj.__class__.__name__
            if (
                ("Conv1D" in p_proj_cls)
                or hasattr(src.c_proj, "nf")
                or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
            ):
                out_w = out_w.t().contiguous()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # ---- Pattern 2: BERT/RoBERTa nested ----
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None
            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # ---- Pattern 3: Qwen/LLaMA ----
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # ---- Pattern 4: DistilBERT ----
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # ---- Fallback: UniversalAdapter ----
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter(old_attn)
                extracted = adapter.extract_qkv_out()
                if extracted:
                    w_q = w_q if w_q is not None else extracted.get("q_weight")
                    b_q = b_q if b_q is not None else extracted.get("q_bias")
                    w_k = w_k if w_k is not None else extracted.get("k_weight")
                    b_k = b_k if b_k is not None else extracted.get("k_bias")
                    w_v = w_v if w_v is not None else extracted.get("v_weight")
                    b_v = b_v if b_v is not None else extracted.get("v_bias")
                    if w_o is None:
                        w_o = extracted.get("out_weight")
                        b_o = extracted.get("out_bias")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        missing = [n for n, w in (("q", w_q), ("k", w_k), ("v", w_v)) if w is None]
        if missing:
            raise RuntimeError(
                f"SqueezingHeadsDistiller.copy_weights: missing mandatory "
                f"projections {missing} from {src_cls_name}."
            )

        # Frobenius checks
        for name, w in (("w_q", w_q), ("w_k", w_k), ("w_v", w_v)):
            if torch.norm(w.float(), p="fro").item() < 1e-4:
                logger.warning("SqueezingHeadsDistiller: %s near-zero Frobenius norm", name)

        # ---- Head importance & selection: SQ2-FIX-02 ----
        importance = self._compute_head_importance(w_q, w_k, w_v)  # (num_heads,)
        num_kept = self.num_kept_heads
        if num_kept < self.num_heads:
            kept_idx, _ = torch.topk(importance, num_kept, largest=True)
            kept_idx, _ = torch.sort(kept_idx)  # stable
        else:
            kept_idx = torch.arange(self.num_heads, device=importance.device)
        self.kept_head_indices = kept_idx.to(torch.long).cpu()

        logger.debug(
            "SqueezingHeadsDistiller: keeping %d/%d heads (indices=%s)",
            num_kept,
            self.num_heads,
            kept_idx.tolist(),
        )

        # ---- Build Q slices for kept heads ----
        D = self.head_dim
        H = self.num_heads
        w_q_heads = w_q[: H * D].view(H, D, -1)  # (H, D, hidden)
        b_q_heads = None
        if b_q is not None and b_q.shape[0] >= H * D:
            b_q_heads = b_q[: H * D].view(H, D)

        kept_q_w = w_q_heads[kept_idx].reshape(num_kept * D, w_q_heads.shape[-1])
        self.q_proj.weight.zero_()
        self.q_proj.weight.copy_(kept_q_w.to(self.q_proj.weight.dtype))

        if b_q_heads is not None and self.q_proj.bias is not None:
            kept_q_b = b_q_heads[kept_idx].reshape(num_kept * D)
            self.q_proj.bias.zero_()
            self.q_proj.bias.copy_(kept_q_b.to(self.q_proj.bias.dtype))

        # ---- K, V: copy ALL num_kv_heads (GQA-safe) ----
        self._copy_kv_full(self.k_proj, w_k, b_k)
        self._copy_kv_full(self.v_proj, w_v, b_v)

        # ---- Output projection ----
        # O must map from kept_q_head_size back to hidden.
        if self.is_nested_encoder:
            # Nested encoder bypasses out_proj entirely.
            nn.init.eye_(self.out_proj.weight)
            if self.out_proj.bias is not None:
                nn.init.zeros_(self.out_proj.bias)
        elif w_o is not None:
            # w_o shape: (hidden, num_heads*D). Take kept head columns.
            if w_o.shape == (self.hidden_size, H * D):
                w_o_heads = w_o.view(self.hidden_size, H, D)
                kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                self.out_proj.weight.zero_()
                self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
            elif w_o.shape == (H * D, self.hidden_size):
                # Transposed
                w_o_t = w_o.t().contiguous()
                w_o_heads = w_o_t.view(self.hidden_size, H, D)
                kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                self.out_proj.weight.zero_()
                self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
            else:
                raise RuntimeError(
                    f"SqueezingHeadsDistiller: unexpected w_o shape "
                    f"{tuple(w_o.shape)}; expected "
                    f"({self.hidden_size}, {H * D}) or ({H * D}, {self.hidden_size})."
                )
            if b_o is not None and self.out_proj.bias is not None:
                m = min(self.out_proj.bias.shape[0], b_o.shape[0])
                self.out_proj.bias.zero_()
                self.out_proj.bias[:m].copy_(b_o[:m].to(self.out_proj.bias.dtype))

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, L, _ = 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
    L_kv_new = kv_input.shape[1]

    # ---- Projections ----
    q = (
        self.q_proj(hidden_states)
        .view(B, L, self.num_kept_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    k_new = (
        self.k_proj(kv_input)
        .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v_new = (
        self.v_proj(kv_input)
        .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # ---- RoPE ----
    if position_embeddings is None:
        position_embeddings = kwargs.get("position_embeddings")
    past_k_probe, _, past_len_rope = self._read_past_kv(past_key_value, layer_idx)
    if position_embeddings is not None and not is_cross:
        cos, sin = position_embeddings
        cos = cos[..., past_len_rope : past_len_rope + L, :]
        sin = sin[..., past_len_rope : past_len_rope + L, :]
        q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin)

    # ---- KV cache ----
    past_k, past_v, past_len = self._read_past_kv(past_key_value, layer_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]

    # ---- GQA expansion ----
    # kept Q heads map to KV heads via (kept_idx // num_kv_groups).
    # Here we simply expand KV to num_kv_groups * num_kv_heads, then select
    # the subset corresponding to kept heads.
    if self.num_kv_heads < self.num_heads:
        # Expand KV to full num_heads then gather kept heads.
        k_full = repeat_kv(k_cat, self.num_kv_groups)  # (B, num_heads, N_kv, D)
        v_full = repeat_kv(v_cat, self.num_kv_groups)
        if self.kept_head_indices is not None and self.num_kept_heads < self.num_heads:
            idx = self.kept_head_indices.to(k_full.device)
            k_full = k_full.index_select(1, idx)
            v_full = v_full.index_select(1, idx)
    else:
        k_full = k_cat
        v_full = v_cat

    # ---- Mask ----
    eff_mask = encoder_attention_mask if is_cross else attention_mask
    token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

    scale = 1.0 / math.sqrt(self.head_dim)

    # ---- SDPA fast path ----
    use_sdpa = not output_attentions and head_mask is None
    if use_sdpa:
        q_sdp = q.to(torch.float16) if q.dtype == torch.float32 else q
        k_sdp = k_full.to(torch.float16) if k_full.dtype == torch.float32 else k_full
        v_sdp = v_full.to(torch.float16) if v_full.dtype == torch.float32 else v_full

        attn_bias_sdp = None
        if token_valid is not None:
            attn_bias_sdp = ((token_valid[:, None, None, :].float() - 1.0) * 1e4).to(
                q_sdp.dtype
            )

        is_causal_sdp = self.is_decoder and not is_cross and past_len == 0 and L > 1

        try:
            attn_out = 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,
            )
            attn_out = (
                attn_out.to(hidden_states.dtype)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.kept_q_head_size)
            )
            attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
            attn_probs_ret = None
        except Exception as e:
            logger.debug("SDPA failed (%s); manual path.", e)
            use_sdpa = False

    if not use_sdpa:
        scores = torch.matmul(q.float(), k_full.float().transpose(-2, -1).contiguous()) * scale

        if self.is_decoder and not is_cross:
            if past_len > 0:
                abs_q = (past_len + torch.arange(L, device=device)).unsqueeze(1)
                abs_k = torch.arange(N_kv, device=device).unsqueeze(0)
                causal = abs_k <= abs_q
            else:
                causal = torch.tril(torch.ones(L, N_kv, device=device, dtype=torch.bool))
            scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

        if token_valid is not None:
            scores = scores.masked_fill(~token_valid[:, None, None, :].bool(), -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        probs = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        probs = self.dropout(probs)

        # head_mask: apply to (B, num_kept_heads, 1, 1)
        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
            if hm.shape[1] != self.num_kept_heads and hm.shape[1] == self.num_heads:
                # head_mask is for the FULL head count; select kept heads.
                if self.kept_head_indices is not None:
                    idx = self.kept_head_indices.to(hm.device)
                    hm = hm.index_select(1, idx)
            probs = probs * hm.to(probs.dtype)

        attn_out = (
            torch.matmul(probs, v_full)
            .transpose(1, 2)
            .contiguous()
            .view(B, L, self.kept_q_head_size)
        )
        attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
        attn_probs_ret = probs if output_attentions else None
    else:
        attn_probs_ret = None

    # ---- Output ----
    if self.is_nested_encoder:
        out = attn_out
    else:
        out = self.out_proj(attn_out)

    if out.shape[-1] != hidden_states.shape[-1]:
        raise RuntimeError(
            f"SqueezingHeadsDistiller: output dim {out.shape[-1]} != "
            f"input dim {hidden_states.shape[-1]}."
        )

    return out, attn_probs_ret

Must return (output, attention_weights_or_None).

class SqueezingHeadsDistiller (config: Any | None = None,
old_attn: torch.nn.modules.module.Module | None = None,
rank: int | None = None,
head_keep_ratio: float | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("squeezing_heads_distiller", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("squeezing_heads_distill", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("squeezing_heads", supports_causal=True, supports_gqa=True)
class SqueezingHeadsDistiller(BaseAttention):
    """
    Attention wrapper that keeps only the top-k most important heads
    (based on pretrained Q/K/V Frobenius norms) while preserving GQA structure.

    Effective head count:
        num_kept_heads = round(head_keep_ratio * num_heads)
        q_head_size   = num_kept_heads * head_dim
        out_head_size = num_kept_heads * head_dim

    KV projections preserve all num_kv_heads (GQA). Kept Q heads are grouped
    to matching KV heads via GQA index mapping.

    Parameter reduction:
        Q : num_heads   * head_dim * hidden  →  num_kept_heads * head_dim * hidden
        O : hidden * num_heads * head_dim    →  hidden * num_kept_heads * head_dim
        K/V: unchanged (GQA-compatible)
    """

    supports_gqa: bool = True
    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        old_attn: nn.Module | None = None,
        rank: int | None = None,
        head_keep_ratio: float | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)

        # ---- 1. Dimensions ----
        self.hidden_size = int(_cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)))
        self.num_heads = int(
            _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "num_heads", kwargs.get("num_attention_heads", 12)),
            )
        )

        cfg_hd = _cfg_get(config, "head_dim", None)
        kw_hd = kwargs.get("head_dim")
        if cfg_hd is not None:
            self.head_dim = int(cfg_hd)
        elif kw_hd is not None:
            self.head_dim = int(kw_hd)
        elif old_attn is not None:
            self.head_dim = None
            for attr in ("q_proj", "query", "q_lin"):
                if hasattr(old_attn, attr):
                    qw = getattr(old_attn, attr).weight
                    self.head_dim = qw.shape[0] // max(1, self.num_heads)
                    break
            if self.head_dim is None:
                self.head_dim = self.hidden_size // max(1, self.num_heads)
        else:
            self.head_dim = self.hidden_size // max(1, self.num_heads)

        cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is not None:
            self.num_kv_heads = int(cfg_kv)
        elif kwargs.get("num_kv_heads") is not None:
            self.num_kv_heads = int(kwargs["num_kv_heads"])
        else:
            self.num_kv_heads = self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))

        # ---- 2. Compute num_kept_heads ----
        # SQ2-FIX-01: derive head_keep_ratio from `rank` if not explicitly given.
        if head_keep_ratio is not None:
            hkr = float(head_keep_ratio)
        elif rank is not None and rank > 0:
            # If rank <= num_heads, interpret as "keep this many heads"
            if rank <= self.num_heads:
                hkr = float(rank) / float(self.num_heads)
            else:
                # Otherwise interpret as "target head_dim"
                effective_hd = min(int(rank), self.head_dim)
                hkr = float(effective_hd) / float(self.head_dim)
        else:
            hkr = 1.0
        self.head_keep_ratio = max(1.0 / self.num_heads, min(1.0, hkr))

        self.num_kept_heads = max(
            1, min(self.num_heads, int(round(self.num_heads * self.head_keep_ratio)))
        )
        self.head_keep_ratio = self.num_kept_heads / self.num_heads

        self.kept_q_head_size = self.num_kept_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim
        self.all_head_size = self.kept_q_head_size

        # GQA mapping: which KV head does each kept Q head map to?
        # Kept heads are the top-k by importance; we must remember their original indices.
        self.num_kv_groups = max(1, self.num_heads // self.num_kv_heads)
        self.kept_head_indices: torch.Tensor | None = None  # filled by copy_weights

        # ---- 3. Decoder detection ----
        self.is_decoder = _detect_decoder(config, is_decoder, old_attn)
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False

        self.lora_rank = int(lora_rank if lora_rank is not None else 16)
        self.lora_alpha = float(lora_alpha)

        # ---- 4. Projections ----
        # SQ2-FIX-01: Q and O are sized to num_kept_heads, not num_heads.
        self.q_proj = LoRALinear(
            self.hidden_size,
            self.kept_q_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.kept_q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = float(
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        # ---- 5. Copy weights from source ----
        src_attn = old_attn if old_attn is not None else kwargs.get("pretrained_attn")
        if src_attn is not None:
            self.copy_weights(src_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 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

    # ------------------------------------------------------------------
    # Head importance: SQ2-FIX-02
    # ------------------------------------------------------------------
    def _compute_head_importance(
        self,
        w_q: torch.Tensor | None,
        w_k: torch.Tensor | None,
        w_v: torch.Tensor | None,
    ) -> torch.Tensor:
        """
        Compute per-head importance as the sum of Frobenius norms of the
        corresponding Q, K, V slices.

        Returns: (num_heads,) tensor of importance scores.
        """
        H = self.num_heads
        D = self.head_dim
        device = None
        for w in (w_q, w_k, w_v):
            if w is not None:
                device = w.device
                break
        if device is None:
            device = torch.device("cpu")

        score = torch.zeros(H, device=device, dtype=torch.float32)

        if w_q is not None and w_q.shape[0] >= H * D:
            w_q_heads = w_q[: H * D].view(H, D, -1)
            score += w_q_heads.float().pow(2).sum(dim=(1, 2))

        if w_k is not None and w_k.shape[0] >= self.num_kv_heads * D:
            # K has num_kv_heads; expand each KV head to the Q heads it serves.
            w_k_heads = w_k[: self.num_kv_heads * D].view(self.num_kv_heads, D, -1)
            k_score = w_k_heads.float().pow(2).sum(dim=(1, 2))  # (num_kv_heads,)
            # Repeat each KV head `num_kv_groups` times to align with Q heads.
            rep = self.num_kv_groups
            k_expanded = k_score.repeat_interleave(rep)
            if k_expanded.numel() < H:
                k_expanded = k_expanded.repeat((H + k_expanded.numel() - 1) // k_expanded.numel())[
                    :H
                ]
            else:
                k_expanded = k_expanded[:H]
            score += k_expanded

        if w_v is not None and w_v.shape[0] >= self.num_kv_heads * D:
            w_v_heads = w_v[: self.num_kv_heads * D].view(self.num_kv_heads, D, -1)
            v_score = w_v_heads.float().pow(2).sum(dim=(1, 2))
            rep = self.num_kv_groups
            v_expanded = v_score.repeat_interleave(rep)
            if v_expanded.numel() < H:
                v_expanded = v_expanded.repeat((H + v_expanded.numel() - 1) // v_expanded.numel())[
                    :H
                ]
            else:
                v_expanded = v_expanded[:H]
            score += v_expanded

        return score

    # ------------------------------------------------------------------
    # copy_weights with head selection
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: nn.Module) -> None:
        with torch.no_grad():
            w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_out_lin = any(hasattr(old_attn, a) for a in ("out_lin", "c_proj", "o_proj"))
            if has_out_lin:
                self.is_nested_encoder = False
            elif (
                hasattr(old_attn, "self")
                and hasattr(old_attn.self, "query")
                and hasattr(old_attn, "output")
                and hasattr(old_attn.output, "dense")
            ) or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention"):
                self.is_nested_encoder = True
            else:
                self.is_nested_encoder = False

            # Determine nested container vs flat
            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    dense = old_attn.output.dense
                    w_o = dense.weight.detach()
                    b_o = dense.bias.detach() if dense.bias is not None else None
            else:
                src = old_attn

            # ---- Pattern 1: GPT-2 fused c_attn ----
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                p_cls = src.c_attn.__class__.__name__
                if (
                    ("Conv1D" in p_cls)
                    or hasattr(src.c_attn, "nf")
                    or (
                        c_attn_w.shape[0] == self.hidden_size
                        and c_attn_w.shape[1] != self.hidden_size
                    )
                ):
                    c_attn_w = c_attn_w.t().contiguous()

                H = self.hidden_size
                w_q = c_attn_w[:H, :]
                w_k = c_attn_w[H : 2 * H, :]
                w_v = c_attn_w[2 * H : 3 * H, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:H]
                    b_k = c_attn_b[H : 2 * H]
                    b_v = c_attn_b[2 * H : 3 * H]

                out_w = src.c_proj.weight.detach()
                p_proj_cls = src.c_proj.__class__.__name__
                if (
                    ("Conv1D" in p_proj_cls)
                    or hasattr(src.c_proj, "nf")
                    or (out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size)
                ):
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # ---- Pattern 2: BERT/RoBERTa nested ----
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # ---- Pattern 3: Qwen/LLaMA ----
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # ---- Pattern 4: DistilBERT ----
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # ---- Fallback: UniversalAdapter ----
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter(old_attn)
                    extracted = adapter.extract_qkv_out()
                    if extracted:
                        w_q = w_q if w_q is not None else extracted.get("q_weight")
                        b_q = b_q if b_q is not None else extracted.get("q_bias")
                        w_k = w_k if w_k is not None else extracted.get("k_weight")
                        b_k = b_k if b_k is not None else extracted.get("k_bias")
                        w_v = w_v if w_v is not None else extracted.get("v_weight")
                        b_v = b_v if b_v is not None else extracted.get("v_bias")
                        if w_o is None:
                            w_o = extracted.get("out_weight")
                            b_o = extracted.get("out_bias")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            missing = [n for n, w in (("q", w_q), ("k", w_k), ("v", w_v)) if w is None]
            if missing:
                raise RuntimeError(
                    f"SqueezingHeadsDistiller.copy_weights: missing mandatory "
                    f"projections {missing} from {src_cls_name}."
                )

            # Frobenius checks
            for name, w in (("w_q", w_q), ("w_k", w_k), ("w_v", w_v)):
                if torch.norm(w.float(), p="fro").item() < 1e-4:
                    logger.warning("SqueezingHeadsDistiller: %s near-zero Frobenius norm", name)

            # ---- Head importance & selection: SQ2-FIX-02 ----
            importance = self._compute_head_importance(w_q, w_k, w_v)  # (num_heads,)
            num_kept = self.num_kept_heads
            if num_kept < self.num_heads:
                kept_idx, _ = torch.topk(importance, num_kept, largest=True)
                kept_idx, _ = torch.sort(kept_idx)  # stable
            else:
                kept_idx = torch.arange(self.num_heads, device=importance.device)
            self.kept_head_indices = kept_idx.to(torch.long).cpu()

            logger.debug(
                "SqueezingHeadsDistiller: keeping %d/%d heads (indices=%s)",
                num_kept,
                self.num_heads,
                kept_idx.tolist(),
            )

            # ---- Build Q slices for kept heads ----
            D = self.head_dim
            H = self.num_heads
            w_q_heads = w_q[: H * D].view(H, D, -1)  # (H, D, hidden)
            b_q_heads = None
            if b_q is not None and b_q.shape[0] >= H * D:
                b_q_heads = b_q[: H * D].view(H, D)

            kept_q_w = w_q_heads[kept_idx].reshape(num_kept * D, w_q_heads.shape[-1])
            self.q_proj.weight.zero_()
            self.q_proj.weight.copy_(kept_q_w.to(self.q_proj.weight.dtype))

            if b_q_heads is not None and self.q_proj.bias is not None:
                kept_q_b = b_q_heads[kept_idx].reshape(num_kept * D)
                self.q_proj.bias.zero_()
                self.q_proj.bias.copy_(kept_q_b.to(self.q_proj.bias.dtype))

            # ---- K, V: copy ALL num_kv_heads (GQA-safe) ----
            self._copy_kv_full(self.k_proj, w_k, b_k)
            self._copy_kv_full(self.v_proj, w_v, b_v)

            # ---- Output projection ----
            # O must map from kept_q_head_size back to hidden.
            if self.is_nested_encoder:
                # Nested encoder bypasses out_proj entirely.
                nn.init.eye_(self.out_proj.weight)
                if self.out_proj.bias is not None:
                    nn.init.zeros_(self.out_proj.bias)
            elif w_o is not None:
                # w_o shape: (hidden, num_heads*D). Take kept head columns.
                if w_o.shape == (self.hidden_size, H * D):
                    w_o_heads = w_o.view(self.hidden_size, H, D)
                    kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                    self.out_proj.weight.zero_()
                    self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
                elif w_o.shape == (H * D, self.hidden_size):
                    # Transposed
                    w_o_t = w_o.t().contiguous()
                    w_o_heads = w_o_t.view(self.hidden_size, H, D)
                    kept_o_w = w_o_heads[:, kept_idx, :].reshape(self.hidden_size, num_kept * D)
                    self.out_proj.weight.zero_()
                    self.out_proj.weight.copy_(kept_o_w.to(self.out_proj.weight.dtype))
                else:
                    raise RuntimeError(
                        f"SqueezingHeadsDistiller: unexpected w_o shape "
                        f"{tuple(w_o.shape)}; expected "
                        f"({self.hidden_size}, {H * D}) or ({H * D}, {self.hidden_size})."
                    )
                if b_o is not None and self.out_proj.bias is not None:
                    m = min(self.out_proj.bias.shape[0], b_o.shape[0])
                    self.out_proj.bias.zero_()
                    self.out_proj.bias[:m].copy_(b_o[:m].to(self.out_proj.bias.dtype))

    # ------------------------------------------------------------------
    # K/V copy helper: preserve ALL KV heads
    # ------------------------------------------------------------------
    def _copy_kv_full(self, target: LoRALinear, w: torch.Tensor, b: torch.Tensor | None) -> None:
        """Copy the full KV projection, sized to num_kv_heads * head_dim.

        Handles source having num_heads heads by mean-pooling to num_kv_heads
        (i.e., converting MHA → GQA). Never truncates heads silently.
        """
        target.weight.zero_()
        if b is not None and target.bias is not None:
            target.bias.zero_()

        target_out = target.weight.shape[0]
        target_in = target.weight.shape[1]
        D = self.head_dim

        # Case A: exact shape match
        if w.shape[0] == target_out and w.shape[1] == target_in:
            target.weight.copy_(w.to(target.weight.dtype))
            if b is not None and target.bias is not None and b.shape[0] == target_out:
                target.bias.copy_(b.to(target.bias.dtype))
            return

        # Case B: source has more heads (MHA → GQA pooling)
        src_out = w.shape[0]
        if src_out > target_out and src_out % D == 0 and target_out % D == 0:
            src_heads = src_out // D
            tgt_heads = target_out // D
            if src_heads % tgt_heads == 0:
                rep = src_heads // tgt_heads
                w_heads = w[:, :target_in].reshape(src_heads, D, target_in)
                w_pooled = w_heads.view(tgt_heads, rep, D, target_in).mean(dim=1)
                w_pooled = w_pooled.reshape(target_out, target_in)
                target.weight.copy_(w_pooled.to(target.weight.dtype))
                if b is not None and target.bias is not None:
                    b_heads = b.reshape(src_heads, D)
                    b_pooled = b_heads.view(tgt_heads, rep, D).mean(dim=1).reshape(target_out)
                    target.bias.copy_(b_pooled.to(target.bias.dtype))
                return

        # Case C: shape mismatch we cannot resolve
        raise RuntimeError(
            f"SqueezingHeadsDistiller: cannot map K/V weight from "
            f"{tuple(w.shape)} to target {tuple(target.weight.shape)}."
        )

    # ------------------------------------------------------------------
    # Mask helper
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(attention_mask, B, N_kv, device):
        if attention_mask is None:
            return None
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 4:
            if attention_mask.dtype == torch.bool:
                v = attention_mask
            elif attention_mask.dtype.is_floating_point:
                v = attention_mask > -1e4
            else:
                v = attention_mask > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    # ------------------------------------------------------------------
    # KV-cache helper
    # ------------------------------------------------------------------
    @staticmethod
    def _read_past_kv(past_key_value, layer_idx=0):
        if past_key_value is None:
            return None, None, 0
        if hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
            idx = min(layer_idx, len(past_key_value.key_cache) - 1)
            k = past_key_value.key_cache[idx]
            v = past_key_value.value_cache[idx]
            return k, v, k.shape[-2]
        if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
            k, v = past_key_value[0], past_key_value[1]
            return k, v, k.shape[-2]
        return None, None, 0

    # ------------------------------------------------------------------
    # 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, L, _ = 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
        L_kv_new = kv_input.shape[1]

        # ---- Projections ----
        q = (
            self.q_proj(hidden_states)
            .view(B, L, self.num_kept_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k_new = (
            self.k_proj(kv_input)
            .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v_new = (
            self.v_proj(kv_input)
            .view(B, L_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- RoPE ----
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        past_k_probe, _, past_len_rope = self._read_past_kv(past_key_value, layer_idx)
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            cos = cos[..., past_len_rope : past_len_rope + L, :]
            sin = sin[..., past_len_rope : past_len_rope + L, :]
            q, k_new = apply_rotary_pos_emb(q, k_new, cos, sin)

        # ---- KV cache ----
        past_k, past_v, past_len = self._read_past_kv(past_key_value, layer_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]

        # ---- GQA expansion ----
        # kept Q heads map to KV heads via (kept_idx // num_kv_groups).
        # Here we simply expand KV to num_kv_groups * num_kv_heads, then select
        # the subset corresponding to kept heads.
        if self.num_kv_heads < self.num_heads:
            # Expand KV to full num_heads then gather kept heads.
            k_full = repeat_kv(k_cat, self.num_kv_groups)  # (B, num_heads, N_kv, D)
            v_full = repeat_kv(v_cat, self.num_kv_groups)
            if self.kept_head_indices is not None and self.num_kept_heads < self.num_heads:
                idx = self.kept_head_indices.to(k_full.device)
                k_full = k_full.index_select(1, idx)
                v_full = v_full.index_select(1, idx)
        else:
            k_full = k_cat
            v_full = v_cat

        # ---- Mask ----
        eff_mask = encoder_attention_mask if is_cross else attention_mask
        token_valid = self._prepare_valid_mask(eff_mask, B, N_kv, device)

        scale = 1.0 / math.sqrt(self.head_dim)

        # ---- SDPA fast path ----
        use_sdpa = not output_attentions and head_mask is None
        if use_sdpa:
            q_sdp = q.to(torch.float16) if q.dtype == torch.float32 else q
            k_sdp = k_full.to(torch.float16) if k_full.dtype == torch.float32 else k_full
            v_sdp = v_full.to(torch.float16) if v_full.dtype == torch.float32 else v_full

            attn_bias_sdp = None
            if token_valid is not None:
                attn_bias_sdp = ((token_valid[:, None, None, :].float() - 1.0) * 1e4).to(
                    q_sdp.dtype
                )

            is_causal_sdp = self.is_decoder and not is_cross and past_len == 0 and L > 1

            try:
                attn_out = 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,
                )
                attn_out = (
                    attn_out.to(hidden_states.dtype)
                    .transpose(1, 2)
                    .contiguous()
                    .view(B, L, self.kept_q_head_size)
                )
                attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
                attn_probs_ret = None
            except Exception as e:
                logger.debug("SDPA failed (%s); manual path.", e)
                use_sdpa = False

        if not use_sdpa:
            scores = torch.matmul(q.float(), k_full.float().transpose(-2, -1).contiguous()) * scale

            if self.is_decoder and not is_cross:
                if past_len > 0:
                    abs_q = (past_len + torch.arange(L, device=device)).unsqueeze(1)
                    abs_k = torch.arange(N_kv, device=device).unsqueeze(0)
                    causal = abs_k <= abs_q
                else:
                    causal = torch.tril(torch.ones(L, N_kv, device=device, dtype=torch.bool))
                scores = scores.masked_fill(~causal.unsqueeze(0).unsqueeze(0), -1e4)

            if token_valid is not None:
                scores = scores.masked_fill(~token_valid[:, None, None, :].bool(), -1e4)

            scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
            probs = F.softmax(scores, dim=-1).to(hidden_states.dtype)
            probs = self.dropout(probs)

            # head_mask: apply to (B, num_kept_heads, 1, 1)
            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
                if hm.shape[1] != self.num_kept_heads and hm.shape[1] == self.num_heads:
                    # head_mask is for the FULL head count; select kept heads.
                    if self.kept_head_indices is not None:
                        idx = self.kept_head_indices.to(hm.device)
                        hm = hm.index_select(1, idx)
                probs = probs * hm.to(probs.dtype)

            attn_out = (
                torch.matmul(probs, v_full)
                .transpose(1, 2)
                .contiguous()
                .view(B, L, self.kept_q_head_size)
            )
            attn_out = torch.nan_to_num(attn_out, nan=0.0, posinf=1e4, neginf=-1e4)
            attn_probs_ret = probs if output_attentions else None
        else:
            attn_probs_ret = None

        # ---- Output ----
        if self.is_nested_encoder:
            out = attn_out
        else:
            out = self.out_proj(attn_out)

        if out.shape[-1] != hidden_states.shape[-1]:
            raise RuntimeError(
                f"SqueezingHeadsDistiller: output dim {out.shape[-1]} != "
                f"input dim {hidden_states.shape[-1]}."
            )

        return out, attn_probs_ret

Attention wrapper that keeps only the top-k most important heads (based on pretrained Q/K/V Frobenius norms) while preserving GQA structure.

Effective head count: num_kept_heads = round(head_keep_ratio * num_heads) q_head_size = num_kept_heads * head_dim out_head_size = num_kept_heads * head_dim

KV projections preserve all num_kv_heads (GQA). Kept Q heads are grouped to matching KV heads via GQA index mapping.

Parameter reduction: Q : num_heads * head_dim * hidden → num_kept_heads * head_dim * hidden O : hidden * num_heads * head_dim → hidden * num_kept_heads * head_dim K/V: unchanged (GQA-compatible)

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.

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

class StableKernelFeatureMap (config)

Expand source code

class StableKernelFeatureMap(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.cfg = config
        self.head_dim = config.head_dim
        self.m = config.num_features
        self.scale = (
            config.scaling if config.scaling is not None else 1.0 / math.sqrt(self.m + config.eps)
        )
        self.register_buffer("proj", self._build_projection())

    def _build_projection(self):
        if self.cfg.seed is not None:
            torch.manual_seed(self.cfg.seed)
        d, m = self.head_dim, self.m
        if self.cfg.method != "orthogonal":
            return torch.randn(d, m) * math.sqrt(d)
        if m <= d:
            q, _ = torch.linalg.qr(torch.randn(d, m))
            return q * math.sqrt(d)
        num_blocks = (m + d - 1) // d
        blocks = [torch.linalg.qr(torch.randn(d, d))[0] for _ in range(num_blocks)]
        return torch.cat(blocks, dim=1)[:, :m] * math.sqrt(d)

    def forward(self, x):
        """FAVOR+ feature map with GLOBAL max subtraction for FP16 stability."""
        with torch.autocast(device_type=x.device.type, enabled=False):
            x_f = x.float()
            phi = torch.matmul(x_f, self.proj.float())
            norm_sq = (x_f * x_f).sum(dim=-1, keepdim=True) / 2.0
            raw = phi - norm_sq

            # FIX: Global max (scalar) subtraction to avoid FP16 overflow.
            # This scalar cancels in the softmax normalization (safe).
            max_val = raw.detach().max()

            exp_in = torch.clamp(
                raw - max_val,
                min=self.cfg.clamp_min,
                max=self.cfg.clamp_max,
            )
            return (self.scale * torch.exp(exp_in)).to(x.dtype)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x) ‑> Callable[..., typing.Any]

Expand source code

def forward(self, x):
    """FAVOR+ feature map with GLOBAL max subtraction for FP16 stability."""
    with torch.autocast(device_type=x.device.type, enabled=False):
        x_f = x.float()
        phi = torch.matmul(x_f, self.proj.float())
        norm_sq = (x_f * x_f).sum(dim=-1, keepdim=True) / 2.0
        raw = phi - norm_sq

        # FIX: Global max (scalar) subtraction to avoid FP16 overflow.
        # This scalar cancels in the softmax normalization (safe).
        max_val = raw.detach().max()

        exp_in = torch.clamp(
            raw - max_val,
            min=self.cfg.clamp_min,
            max=self.cfg.clamp_max,
        )
        return (self.scale * torch.exp(exp_in)).to(x.dtype)

FAVOR+ feature map with GLOBAL max subtraction for FP16 stability.

class TLinFormerAttention (config: Any | None = None,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("tlinformer")
@ATTENTION_REGISTRY.register("tlinformer_attention")
class TLinFormerAttention(BaseAttention):
    """
    TLinFormer: linear-time attention with sequence-dimension low-rank projection.

    For an encoder with N tokens and rank r ≤ N:
        K' = E K,  V' = F V,  E ∈ R^{r×N},  F ∈ R^{r×N}
        scores = Q K'ᵀ / √d      # (B, H, N, r)
        O = softmax(scores) V'   # (B, H, N, d)

    Complexity O(N · r · d) instead of O(N² · d).

    For decoders (causal), we deliberately fall back to dense attention. A
    projected key aggregates information from all N positions, so a causal
    mask over the projected keys cannot enforce causality — the model would
    silently leak future information. Dense causal attention is correct.
    """

    # TLF-FIX-04
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = 64,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank or 64)
        self.config = config
        comp_config = getattr(config, "compression_config", None)

        # ---- hidden size / heads ----
        self.hidden_size = getattr(config, "hidden_size", kwargs.get("hidden_size", 768))
        if config is None or not hasattr(config, "hidden_size"):
            if pretrained_attn is not None:
                for cand in ("q_lin", "query", "q_proj", "c_attn"):
                    mod = getattr(pretrained_attn, cand, None)
                    if mod is not None and hasattr(mod, "weight"):
                        w = mod.weight
                        if w.ndim == 2:
                            self.hidden_size = w.shape[1]
                            break

        self.num_heads = int(
            getattr(config, "num_attention_heads", kwargs.get("num_attention_heads", 12))
            if config
            else 12
        )
        self.num_kv_heads = int(
            getattr(
                config,
                "num_key_value_heads",
                kwargs.get("num_key_value_heads", self.num_heads),
            )
            if config
            else self.num_heads
        )
        self.head_dim = int(
            getattr(config, "head_dim", None)
            or kwargs.get("head_dim")
            or (self.hidden_size // max(1, self.num_heads))
        )
        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))

        # ---- rank: TLF-FIX-13 clamp to head_dim ----
        requested_rank = int(
            _cfg_get(
                comp_config,
                "tlinformer_rank",
                getattr(config, "tlinformer_rank", None) if config else None,
            )
            or kwargs.pop("rank", rank or 64)
        )
        self.rank = max(1, min(requested_rank, self.head_dim))

        self.max_seq_len = int(
            getattr(config, "max_position_embeddings", kwargs.get("max_seq_len", 2048))
            if config
            else 2048
        )

        # ---- decoder detection ----
        _mt = (getattr(config, "model_type", "") if config else "").lower()
        is_dec = bool(
            getattr(config, "is_decoder", kwargs.get("is_decoder", False)) if config else False
        )
        if any(
            _mt.startswith(p)
            for p in (
                "gpt",
                "qwen",
                "llama",
                "mistral",
                "gemma",
                "phi",
                "opt",
                "bloom",
                "falcon",
                "olmo",
                "mpt",
                "stablelm",
            )
        ):
            is_dec = True
        if kwargs.get("is_causal", False):
            is_dec = True
        self.is_decoder = is_dec
        self.is_causal = is_dec
        self.is_nested_encoder = False

        # ---- LoRA ----
        self.lora_rank = int(
            getattr(
                config,
                "lora_rank",
                getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)),
            )
            if config
            else 16
        )
        self.lora_alpha = float(
            getattr(
                config,
                "lora_alpha",
                getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)),
            )
            if config
            else 32.0
        )
        self.freeze_base = self.lora_rank > 0

        # ---- projections ----
        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=self.freeze_base,
        )
        self.k_proj = LoRALinear(
            self.hidden_size,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        self.v_proj = LoRALinear(
            self.hidden_size,
            self.all_kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=self.freeze_base,
        )
        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=self.freeze_base,
        )

        # ---- sequence projections (TLF-FIX-06: buffers, not params) ----
        # Shape (rank, max_seq_len):  E @ K projects N → rank
        proj_k = torch.empty(self.rank, self.max_seq_len)
        proj_v = torch.empty(self.rank, self.max_seq_len)
        nn.init.normal_(proj_k, std=0.02)
        nn.init.normal_(proj_v, std=0.02)
        # Identity init on the leading diagonal for the "rank = seq" boundary
        md = min(self.rank, self.max_seq_len)
        proj_k[:md, :md].fill_diagonal_(1.0)
        proj_v[:md, :md].fill_diagonal_(1.0)
        self.register_buffer("proj_k", proj_k, persistent=True)
        self.register_buffer("proj_v", proj_v, persistent=True)

        dropout_p = float(
            getattr(config, "attention_probs_dropout_prob", kwargs.get("dropout", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        src = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn")
        if src is not None:
            self.copy_weights(src)

    # ---- compatibility 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 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

    @property
    def E(self):
        return self.proj_k

    @property
    def F(self):
        return self.proj_v

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, src: Any) -> None:
        if src is None:
            return

        cls_name = src.__class__.__name__
        has_out_lin = any(hasattr(src, a) for a in ("out_lin", "c_proj", "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 = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # ---- GPT-2 fused c_attn / c_proj ----
        if hasattr(src, "c_attn") and src.c_attn is not None:
            c = src.c_attn
            base_c = getattr(c, "base_layer", c)
            cw = getattr(base_c, "weight", getattr(c, "weight", None))
            if cw is not None:
                cw = cw.data.clone()
                is_conv = (
                    "Conv1D" in c.__class__.__name__
                    or "Conv1D" in base_c.__class__.__name__
                    or hasattr(c, "nf")
                    or hasattr(base_c, "nf")
                    or (cw.ndim == 2 and cw.shape[0] < cw.shape[1])
                )
                if is_conv:
                    cw = cw.t().contiguous()
                split = cw.size(0) // 3
                bmod = getattr(base_c, "bias", getattr(c, "bias", None))
                cb = bmod.data.clone() if bmod is not None else None
                extracted["q"] = (cw[:split], cb[:split] if cb is not None else None)
                extracted["k"] = (
                    cw[split : 2 * split],
                    cb[split : 2 * split] if cb is not None else None,
                )
                extracted["v"] = (
                    cw[2 * split :],
                    cb[2 * split :] if cb is not None else None,
                )

            if hasattr(src, "c_proj") and src.c_proj is not None:
                o = src.c_proj
                base_o = getattr(o, "base_layer", o)
                ow = getattr(base_o, "weight", getattr(o, "weight", None))
                if ow is not None:
                    ow = ow.data.clone()
                    is_conv_o = (
                        "Conv1D" in o.__class__.__name__
                        or "Conv1D" in base_o.__class__.__name__
                        or hasattr(o, "nf")
                        or hasattr(base_o, "nf")
                        or (ow.ndim == 2 and ow.shape[0] < ow.shape[1])
                    )
                    if is_conv_o:
                        ow = ow.t().contiguous()
                    bo = getattr(base_o, "bias", getattr(o, "bias", None))
                    extracted["o"] = (ow, bo.data.clone() if bo is not None else None)

        # ---- RoBERTa / BERT nested ----
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            for key, sub in (("q", "query"), ("k", "key"), ("v", "value")):
                m = getattr(s, sub, None)
                if m is not None and hasattr(m, "weight"):
                    extracted[key] = (
                        m.weight.data.clone(),
                        m.bias.data.clone() if m.bias is not None else None,
                    )
            if hasattr(src, "output") and hasattr(src.output, "dense"):
                om = src.output.dense
                extracted["o"] = (
                    om.weight.data.clone(),
                    om.bias.data.clone() if om.bias is not None else None,
                )

        # ---- flat modules ----
        else:
            cand = {
                "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 key, names in cand.items():
                for n in names:
                    if hasattr(src, n):
                        m = getattr(src, n)
                        if m is not None and hasattr(m, "weight"):
                            extracted[key] = (
                                m.weight.data.clone(),
                                m.bias.data.clone() if m.bias is not None else None,
                            )
                            break

        # ---- UniversalAdapter fallback ----
        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

                ua = UniversalAdapter()
                for k in missing:
                    w, b = ua.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"TLinFormer copy_weights: missing {missing_after} from {src.__class__.__name__}."
            )

        target_map = {
            "q": self.q_proj,
            "k": self.k_proj,
            "v": self.v_proj,
            "o": self.out_proj,
        }

        with torch.no_grad():
            for key, target in target_map.items():
                w, b = extracted[key]
                if w is None:
                    if key == "o" and self.is_nested_encoder:
                        nn.init.eye_(target.weight)
                        if target.bias is not None:
                            nn.init.zeros_(target.bias)
                    continue

                if w.ndim != 2:
                    w = w.reshape(w.size(0), -1)

                # TLF-FIX-09: log transposition
                if (
                    w.shape != target.weight.shape
                    and w.t().contiguous().shape == target.weight.shape
                ):
                    logger.debug(
                        "TLinFormer: transposing '%s' from %s to %s",
                        key,
                        tuple(w.shape),
                        tuple(target.weight.shape),
                    )
                    w = w.t().contiguous()

                if w.shape != target.weight.shape:
                    logger.warning(
                        "TLinFormer: '%s' shape mismatch %s vs target %s; truncating.",
                        key,
                        tuple(w.shape),
                        tuple(target.weight.shape),
                    )

                out_d = min(target.weight.size(0), w.size(0))
                in_d = min(target.weight.size(1), w.size(1))
                target.weight.data[:out_d, :in_d].copy_(
                    w[:out_d, :in_d].to(target.weight.device, target.weight.dtype)
                )

                if b is not None and target.bias is not None:
                    mn = min(target.bias.numel(), b.numel())
                    target.bias.data[:mn].copy_(b[:mn].to(target.bias.device, target.bias.dtype))
                elif target.bias is not None:
                    nn.init.zeros_(target.bias)

                fnorm = torch.linalg.norm(target.weight.data.float()).item()
                if fnorm < 1e-4:
                    logger.warning("TLinFormer: '%s' near-zero norm after copy", key)

        # ---- LoRA copy (TLF-FIX-11) ----
        for key, target in target_map.items():
            src_mod = None
            if key == "o" and hasattr(src, "output") and hasattr(src.output, "dense"):
                src_mod = src.output.dense
            elif hasattr(src, "self"):
                sub_map = {"q": "query", "k": "key", "v": "value"}
                src_mod = getattr(src.self, sub_map.get(key, ""), None)
            else:
                for n in (
                    "q_lin",
                    "q_proj",
                    "query",
                    "k_lin",
                    "k_proj",
                    "key",
                    "v_lin",
                    "v_proj",
                    "value",
                    "out_lin",
                    "out_proj",
                    "dense",
                    "o_proj",
                ):
                    if hasattr(src, n):
                        cand = getattr(src, n)
                        if key in n:
                            src_mod = cand
                            break
            if src_mod is not None:
                sA = getattr(src_mod, "lora_A", None)
                sB = getattr(src_mod, "lora_B", None)
                if (
                    sA is not None
                    and sB is not None
                    and target.lora_A is not None
                    and target.lora_B is not None
                ):
                    wA = getattr(sA, "weight", sA)
                    wB = getattr(sB, "weight", sB)
                    try:
                        if wA.shape == target.lora_A.shape:
                            target.lora_A.data.copy_(wA.data.to(target.lora_A.dtype))
                        if wB.shape == target.lora_B.shape:
                            target.lora_B.data.copy_(wB.data.to(target.lora_B.dtype))
                    except Exception as e:
                        logger.debug("TLinFormer LoRA copy failed for '%s': %s", key, e)

        if self.freeze_base:
            for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
                m.freeze_base_weights()

    # ------------------------------------------------------------------
    # helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _prepare_valid_mask(attention_mask, B, N, device):
        if attention_mask is None:
            return torch.ones(B, N, device=device, dtype=torch.bool)
        if attention_mask.dim() == 2:
            if attention_mask.dtype == torch.bool:
                return attention_mask
            if attention_mask.dtype.is_floating_point:
                return attention_mask > -1e4
            return attention_mask > 0
        if attention_mask.dim() == 3:
            m = attention_mask
            if m.dtype == torch.bool:
                return m.squeeze(1)
            if m.dtype.is_floating_point:
                return (m > -1e4).squeeze(1)
            return (m > 0).squeeze(1)
        if attention_mask.dim() == 4:
            m = attention_mask
            if m.dtype == torch.bool:
                v = m
            elif m.dtype.is_floating_point:
                v = m > -1e4
            else:
                v = m > 0
            return v.any(dim=-2).squeeze(1)
        raise ValueError(f"Unsupported attention_mask dim: {attention_mask.dim()}")

    def _projection_matrix(self, name: str, L: int, dtype, device) -> torch.Tensor:
        """
        TLF-FIX-07 / TLF-FIX-08: return a projection matrix of shape (r, L).

        For L ≤ max_seq_len, we slice `proj_[k|v][:, :L]`.
        For L > max_seq_len, we log a warning and clip to max_seq_len (no
        interpolate hack).
        """
        buf = getattr(self, name)  # (rank, max_seq_len)
        buf.shape[0]
        if self.max_seq_len < L:
            logger.warning(
                "TLinFormer: sequence length %d exceeds max_seq_len %d; "
                "projecting with clipped columns (results may degrade).",
                L,
                self.max_seq_len,
            )
            L_eff = self.max_seq_len
        else:
            L_eff = L
        P = buf[:, :L_eff].to(dtype=dtype, device=device)  # (r, L_eff)
        if L_eff < L:
            # Pad with zeros so matmul over the full sequence still works
            P = F.pad(P, (0, L - L_eff), value=0.0)
        return P

    # ------------------------------------------------------------------
    # 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,
        *args,
        **kwargs,
    ):
        B, N, _ = hidden_states.shape
        H, Dh = self.num_heads, self.head_dim

        # ---- projections ----
        q = self.q_proj(hidden_states).view(B, N, H, Dh).transpose(1, 2).contiguous()
        k = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_kv_heads, Dh)
            .transpose(1, 2)
            .contiguous()
        )

        # ---- RoPE ----
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # ---- GQA ----
        if self.kv_repeat > 1:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        # ---- valid mask ----
        valid = self._prepare_valid_mask(attention_mask, B, N, hidden_states.device)  # (B, N)
        key_mask = valid.unsqueeze(1).unsqueeze(2)  # (B,1,1,N)
        q_mask = valid.unsqueeze(1).unsqueeze(-1)  # (B,1,N,1)

        scale = 1.0 / math.sqrt(Dh)

        # ---- decoder / causal path: dense attention (TLF-FIX-01) ----
        if self.is_decoder:
            scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
            causal = torch.triu(
                torch.full((N, N), -1e4, device=scores.device, dtype=torch.float32),
                diagonal=1,
            )
            scores = scores + causal.unsqueeze(0).unsqueeze(0)
            # TLF-FIX-02 / TLF-FIX-10: mask KEYS, not queries
            scores = scores.masked_fill(~key_mask, -1e4)
            attn = F.softmax(scores, dim=-1).to(hidden_states.dtype)
            attn = torch.nan_to_num(attn, nan=0.0)
            attn = self.dropout(attn)
            # TLF-FIX-03: head_mask correctly
            if head_mask is not None:
                attn = attn * self._reshape_head_mask(head_mask, B, H, attn.device, attn.dtype)
            ctx = torch.matmul(attn, v)
            # zero-out padded query positions
            ctx = ctx * q_mask.to(ctx.dtype)
            ctx = ctx.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
            out = ctx if self.is_nested_encoder else self.out_proj(ctx)
            return out, (attn if output_attentions else None)

        # ---- encoder: linear-time LinFormer ----
        # TLF-FIX-07: if N ≤ rank, projection degenerates to identity — same
        # as the base model. Otherwise use the low-rank projection.
        if self.rank >= N:
            # Degenerate case → exact dense attention (safe baseline)
            scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
            # TLF-FIX-10: mask KEYS
            scores = scores.masked_fill(~key_mask, -1e4)
            attn = F.softmax(scores, dim=-1).to(hidden_states.dtype)
            attn = torch.nan_to_num(attn, nan=0.0)
            attn = self.dropout(attn)
            if head_mask is not None:
                attn = attn * self._reshape_head_mask(head_mask, B, H, attn.device, attn.dtype)
            ctx = torch.matmul(attn, v)
            ctx = ctx * q_mask.to(ctx.dtype)
            ctx = ctx.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
            out = ctx if self.is_nested_encoder else self.out_proj(ctx)
            return out, (attn if output_attentions else None)

        # Proper LinFormer: N > rank
        P_k = self._projection_matrix(
            "proj_k", N, hidden_states.dtype, hidden_states.device
        )  # (r, N)
        P_v = self._projection_matrix(
            "proj_v", N, hidden_states.dtype, hidden_states.device
        )  # (r, N)

        # Mask padded keys before projecting: k_masked = k * key_mask
        k_for_proj = k * key_mask.to(k.dtype)  # (B, H, N, Dh)
        v_for_proj = v * key_mask.to(v.dtype)

        # Project sequence dimension: (r, N) @ (B, H, N, Dh) → (B, H, r, Dh)
        K_proj = torch.matmul(P_k, k_for_proj)  # (B, H, r, Dh)
        V_proj = torch.matmul(P_v, v_for_proj)  # (B, H, r, Dh)

        # scores: (B, H, N, r)
        scores = torch.matmul(q.float(), K_proj.float().transpose(-2, -1).contiguous()) * scale

        # TLF-FIX-10: mask keys (padded positions were zeroed but we still
        # want to explicitly prevent attention to them)
        # Since padded keys were zeroed BEFORE projection, K_proj for those
        # positions is 0 → their contribution to scores is 0. We don't need
        # to mask them again here. But if attention_mask carried NaN or other
        # signals, we would still mask. For consistency, no-op here.

        attn = F.softmax(scores, dim=-1).to(hidden_states.dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        # TLF-FIX-03: head_mask correctly
        if head_mask is not None:
            attn = attn * self._reshape_head_mask(head_mask, B, H, attn.device, attn.dtype)

        # context: (B, H, N, r) @ (B, H, r, Dh) → (B, H, N, Dh)
        ctx = torch.matmul(attn, V_proj)
        ctx = ctx * q_mask.to(ctx.dtype)

        ctx = ctx.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        out = ctx if self.is_nested_encoder else self.out_proj(ctx)
        return out, (attn if output_attentions else None)

    # ------------------------------------------------------------------
    @staticmethod
    def _reshape_head_mask(head_mask, B, H, device, dtype):
        """TLF-FIX-03: correct shape handling for head_mask."""
        if head_mask is None:
            return None
        hm = head_mask.to(device=device, dtype=dtype)
        if hm.dim() == 1:
            # (H,) → (1, H, 1, 1)
            return hm.view(1, H, 1, 1)
        if hm.dim() == 2:
            # (B, H) → (B, H, 1, 1)
            return hm.view(hm.shape[0], H, 1, 1)
        if hm.dim() == 3:
            # (B, H, 1) or (num_layers, B, H)? Standard: (num_layers, H, 1)
            # We interpret as (B, H, 1) → (B, H, 1, 1)
            return hm.unsqueeze(-1)
        if hm.dim() == 4:
            return hm
        raise ValueError(f"Unsupported head_mask dim: {hm.dim()}")

TLinFormer: linear-time attention with sequence-dimension low-rank projection.

For an encoder with N tokens and rank r ≤ N: K' = E K, V' = F V, E ∈ R^{r×N}, F ∈ R^{r×N} scores = Q K'ᵀ / √d # (B, H, N, r) O = softmax(scores) V' # (B, H, N, d)

Complexity O(N · r · d) instead of O(N² · d).

For decoders (causal), we deliberately fall back to dense attention. A projected key aggregates information from all N positions, so a causal mask over the projected keys cannot enforce causality — the model would silently leak future information. Dense causal attention is correct.

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.

Instance variables

prop E

Expand source code

@property
def E(self):
    return self.proj_k

prop F

Expand source code

@property
def F(self):
    return self.proj_v

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

class TPA (config: Any | None = None,
rank: int = 8,
pretrained_attn: nn.Module | None = None,
lora_rank: int = 0,
lora_alpha: float = 1.0,
is_decoder: bool | None = None,
max_cp_error: float = 0.35,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("tpa", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("tpa_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("tp_attention", supports_causal=True, supports_gqa=True)
class TPAttention(BaseAttention):
    supports_gqa: bool = True
    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 8,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        is_decoder: bool | None = None,
        max_cp_error: float = 0.35,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        src = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")
        self.max_cp_error = float(max_cp_error)

        # ---- hidden size ----
        dim = None
        if config is not None and not isinstance(config, int):
            dim = _cfg_get(
                config,
                "hidden_size",
                _cfg_get(config, "dim", _cfg_get(config, "d_model", None)),
            )
        if dim is None and src is not None:
            dim = getattr(src, "hidden_size", getattr(src, "embed_dim", None))
            if dim is None:
                for attr in ("q_proj", "q_lin", "query"):
                    m = getattr(src, attr, None)
                    if m is not None and hasattr(m, "in_features"):
                        dim = m.in_features
                        break
            if dim is None and hasattr(src, "c_attn") and hasattr(src.c_attn, "weight"):
                w = src.c_attn.weight
                dim = w.shape[0] if w.shape[0] != 3 * w.shape[1] else w.shape[1]
        if dim is None:
            dim = kwargs.get("hidden_size", 768)

        # ---- heads ----
        num_heads = None
        if config is not None and not isinstance(config, int):
            num_heads = _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "n_heads", _cfg_get(config, "num_heads", None)),
            )
        if num_heads is None and src is not None:
            num_heads = getattr(
                src,
                "num_heads",
                getattr(src, "num_attention_heads", getattr(src, "n_head", None)),
            )
        if num_heads is None:
            num_heads = 12

        self.dim = int(dim)
        self.num_heads = max(1, int(num_heads))

        # ---- head_dim ----
        cfg_hd = _cfg_get(config, "head_dim", None) if config is not None else None
        kw_hd = kwargs.get("head_dim")
        if cfg_hd is not None:
            self.head_dim = int(cfg_hd)
        elif kw_hd is not None:
            self.head_dim = int(kw_hd)
        elif src is not None and getattr(src, "head_dim", None) is not None:
            self.head_dim = int(src.head_dim)
        else:
            self.head_dim = self.dim // self.num_heads

        # ---- GQA ----
        cfg_kv = (
            _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
            if config is not None
            else None
        )
        if cfg_kv is None and src is not None:
            cfg_kv = getattr(src, "num_key_value_heads", getattr(src, "num_kv_heads", None))
        self.num_kv_heads = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_kv_groups = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.q_head_size = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim

        # ---- rank (head-rank for tensor product) ----
        raw_r = kwargs.pop("rank", _cfg_get(config, "rank", rank))
        # TPA-FIX-02: rank must satisfy r*r == head_dim for tensor product.
        # If not, we use the qk_proj path (proper projection).
        self.rank = int(raw_r)
        self._use_tensor_product = self.rank * self.rank == self.head_dim
        # Effective qk dimension after tensor-product or qk_proj
        self.effective_qk_dim = self.rank * self.rank if self._use_tensor_product else self.head_dim

        self.lora_rank = int(kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank)))
        self.lora_alpha = float(
            kwargs.get("lora_alpha", _cfg_get(config, "lora_alpha", lora_alpha))
        )

        # ---- decoder ----
        _mt = str(_cfg_get(config, "model_type", "") or "").lower() if config is not None else ""
        _dec_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        elif config is not None and (
            _cfg_get(config, "is_decoder", False)
            or _cfg_get(config, "is_causal", False)
            or any(_mt.startswith(p) for p in _dec_prefixes)
        ):
            self.is_decoder = True
        else:
            self.is_decoder = bool(kwargs.get("is_decoder", False))

        self.is_nested_encoder = False
        # TPA-FIX-14: gradient checkpointing only when explicitly requested.
        self.gradient_checkpointing = bool(kwargs.get("gradient_checkpointing", False))

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and src is not None:
            p = next(src.parameters(), None)
            if p is not None:
                target_dtype = p.dtype

        # ---- projections ----
        # TPA-FIX-11: LoRA is applied to a separate low-rank residual, not to
        # the tensor-product factors. This keeps the semantics well-defined.
        self.q_proj1 = LoRALinear(
            self.dim,
            self.num_heads * self.rank,
            bias=False,
            lora_rank=0,
            lora_alpha=1.0,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.q_proj2 = LoRALinear(
            self.dim,
            self.num_heads * self.rank,
            bias=False,
            lora_rank=0,
            lora_alpha=1.0,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.k_proj1 = LoRALinear(
            self.dim,
            self.num_kv_heads * self.rank,
            bias=False,
            lora_rank=0,
            lora_alpha=1.0,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.k_proj2 = LoRALinear(
            self.dim,
            self.num_kv_heads * self.rank,
            bias=False,
            lora_rank=0,
            lora_alpha=1.0,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.out_proj = LoRALinear(
            self.q_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )

        # TPA-FIX-08/15: when not using tensor-product, use a proper projection
        # from (rank*rank) → head_dim. Otherwise the tensor product already
        # produces head_dim directly.
        if not self._use_tensor_product:
            n_in = self.rank * self.rank
            n_out = self.head_dim
            self.qk_proj = nn.Parameter(torch.zeros(n_out, n_in), requires_grad=False)
        else:
            self.register_parameter("qk_proj", None)

        # TPA-FIX-11: dedicated LoRA on the effective Q/K projections.
        # We store them as LoRALinear to align with the framework convention.
        # They apply to the concatenated (H*hd) space, not the factorized dims.
        if self.lora_rank > 0:
            self.q_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.dim, dtype=target_dtype))
            self.q_lora_B = nn.Parameter(
                torch.zeros(self.q_head_size, self.lora_rank, dtype=target_dtype)
            )
            self.k_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.dim, dtype=target_dtype))
            self.k_lora_B = nn.Parameter(
                torch.zeros(self.kv_head_size, self.lora_rank, dtype=target_dtype)
            )
            nn.init.kaiming_uniform_(self.q_lora_A, a=math.sqrt(5))
            nn.init.kaiming_uniform_(self.k_lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("q_lora_A", None)
            self.register_parameter("q_lora_B", None)
            self.register_parameter("k_lora_A", None)
            self.register_parameter("k_lora_B", None)

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

        if target_dtype is not None:
            for proj in (
                self.q_proj1,
                self.q_proj2,
                self.k_proj1,
                self.k_proj2,
                self.v_proj,
                self.out_proj,
            ):
                proj.to(target_dtype)

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        with torch.no_grad():
            w_q = w_k = w_v = w_o = None
            b_q = b_k = b_v = b_o = None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_no_out = not any(
                hasattr(old_attn, a)
                for a in ("dense", "out_proj", "o_proj", "c_proj", "out_lin", "proj")
            )
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query") and has_no_out
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight.detach()
                    b_o = (
                        old_attn.output.dense.bias.detach()
                        if old_attn.output.dense.bias is not None
                        else None
                    )
            else:
                src = old_attn

            # GPT-2
            if hasattr(src, "c_attn"):
                W = src.c_attn.weight.detach()
                B_vec = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                if W.shape[0] != 3 * self.dim and W.shape[1] == 3 * self.dim:
                    W = W.t().contiguous()
                if W.shape[0] == 3 * self.dim:
                    w_q, w_k, w_v = W.chunk(3, dim=0)
                    if B_vec is not None:
                        b_q, b_k, b_v = B_vec.chunk(3, dim=0)
                if hasattr(src, "c_proj"):
                    w_o = src.c_proj.weight.detach()
                    if w_o.shape[0] != self.dim and w_o.shape[1] == self.dim:
                        w_o = w_o.t().contiguous()
                    b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # BERT/RoBERTa
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None
                if w_o is None:
                    for a in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, a):
                            proj = getattr(src, a)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Qwen / LLaMA
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # DistilBERT
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    if w_q is None:
                        w_q, b_q = adapter.get_weight_bias(src, "q")
                    if w_k is None:
                        w_k, b_k = adapter.get_weight_bias(src, "k")
                    if w_v is None:
                        w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            missing = []
            if w_q is None:
                missing.append("Q")
            if w_k is None:
                missing.append("K")
            if w_v is None:
                missing.append("V")
            if w_o is None and not self.is_nested_encoder:
                missing.append("O")
            if missing:
                logger.warning(
                    "TPAttention: failed to extract %s from %s",
                    missing,
                    type(old_attn).__name__,
                )
                return

            # ---- Factorize Q and K ----
            def _factorize_qk(W, proj1, proj2, num_h):
                """TPA-FIX-07/15: try CP; if error > max_cp_error, use SVD."""
                W_f = W.float()
                in_dim = W_f.shape[1]
                r_rows = num_h * self.rank

                cp_ok = False
                if self._use_tensor_product:
                    try:
                        A, B = _cp_decompose_qk_fast(
                            W_f,
                            num_h,
                            self.rank,
                            self.head_dim,
                            n_iter=3,
                        )
                        err = _cp_reconstruction_error(
                            W_f,
                            A,
                            B,
                            num_h,
                            self.rank,
                            self.head_dim,
                        )
                        if err <= self.max_cp_error and not (
                            torch.isnan(A).any() or torch.isnan(B).any()
                        ):
                            proj1.weight.data.zero_()
                            proj2.weight.data.zero_()
                            proj1.weight.data[:r_rows, :in_dim].copy_(
                                A[:r_rows, :in_dim].to(proj1.weight.dtype)
                            )
                            proj2.weight.data[:r_rows, :in_dim].copy_(
                                B[:r_rows, :in_dim].to(proj2.weight.dtype)
                            )
                            logger.info("TPAttention: CP OK (err=%.4f) for %d heads", err, num_h)
                            cp_ok = True
                        else:
                            logger.warning(
                                "TPAttention: CP error %.4f > max %.4f; using SVD",
                                err,
                                self.max_cp_error,
                            )
                    except Exception as e:
                        logger.warning("TPAttention: CP failed: %s; using SVD", e)

                if not cp_ok:
                    # Fallback: use SVD factorized projection, and mark qk_proj
                    # to be identity-like (we set proj2 = I to preserve structure).
                    # We store SVD factors into proj1 (all rows) and proj2 as identity.
                    # This is a slight approximation — the SVD maps W → (U√S, Vt)
                    # but the tensor-product forward requires matching shapes.
                    # Instead, we use the qk_proj path.
                    A, B = _svd_factorize(W_f, self.rank * self.rank)
                    # Store full rank factorization in proj1 (as (H*r, in)) and
                    # leave proj2 as orthogonal init. The qk_proj path will
                    # combine them.
                    r_eff = min(A.shape[1], r_rows)
                    proj1.weight.data.zero_()
                    proj2.weight.data.zero_()
                    # Store A directly in proj1 (truncated to r_rows)
                    proj1.weight.data[:r_eff, :in_dim].copy_(
                        A[:r_eff, :in_dim].to(proj1.weight.dtype)
                    )
                    # Store B in proj2 (truncated to r_rows)
                    proj2.weight.data[:r_eff, :in_dim].copy_(
                        B[:r_eff, :in_dim].to(proj2.weight.dtype)
                    )
                    # Set qk_proj to identity-like on the effective dim
                    if self.qk_proj is not None:
                        n = min(self.qk_proj.shape)
                        self.qk_proj.data.zero_()
                        for i in range(n):
                            self.qk_proj.data[i, i] = 1.0
                    logger.info("TPAttention: SVD fallback applied for %d heads", num_h)

                proj1.weight.requires_grad = False
                proj2.weight.requires_grad = False

            _factorize_qk(w_q, self.q_proj1, self.q_proj2, self.num_heads)
            _factorize_qk(w_k, self.k_proj1, self.k_proj2, self.num_kv_heads)

            # ---- Copy V and O directly ----
            def _copy_layer(target, w, b):
                if w is None:
                    return
                if w.shape == target.weight.shape:
                    target.weight.copy_(w.to(target.weight.dtype))
                elif w.t().contiguous().shape == target.weight.shape:
                    target.weight.copy_(w.t().contiguous().to(target.weight.dtype))
                else:
                    m0 = min(target.weight.shape[0], w.shape[0])
                    m1 = min(target.weight.shape[1], w.shape[1])
                    target.weight.zero_()
                    target.weight[:m0, :m1].copy_(w[:m0, :m1].to(target.weight.dtype))
                if b is not None and target.bias is not None:
                    mb = min(target.bias.shape[0], b.shape[0])
                    target.bias.zero_()
                    target.bias[:mb].copy_(b[:mb].to(target.bias.dtype))
                elif target.bias is not None:
                    target.bias.zero_()

            _copy_layer(self.v_proj, w_v, b_v)
            if w_o is not None:
                _copy_layer(self.out_proj, w_o, b_o)

            if hasattr(old_attn, "is_decoder"):
                self.is_decoder = bool(old_attn.is_decoder)
            elif hasattr(old_attn, "c_attn") or any(
                x in old_attn.__class__.__name__.lower()
                for x in ("gpt", "qwen", "llama", "decoder")
            ):
                self.is_decoder = True

    # ------------------------------------------------------------------
    # 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,
        *args,
        **kwargs,
    ):
        B, L, D = hidden_states.shape
        H, H_kv, r, hd = self.num_heads, self.num_kv_heads, self.rank, self.head_dim

        # ---- Q, K via tensor product ----
        def _compute_q(x):
            q1 = self.q_proj1(x).view(B, L, H, r)
            q2 = self.q_proj2(x).view(B, L, H, r)
            out_q = (q1.unsqueeze(-1) * q2.unsqueeze(-2)).flatten(-2, -1)  # (B, L, H, r*r)
            if self.qk_proj is not None:
                out_q = F.linear(out_q, self.qk_proj.to(out_q.dtype))
            # TPA-FIX-11: add LoRA residual on the (H*hd) space
            if self.q_lora_A is not None:
                lora = F.linear(F.linear(x, self.q_lora_A), self.q_lora_B)
                out_q = out_q + lora.view(B, L, H, hd) * (self.lora_alpha / max(1, self.lora_rank))
            return out_q

        def _compute_k(x):
            k1 = self.k_proj1(x).view(B, L, H_kv, r)
            k2 = self.k_proj2(x).view(B, L, H_kv, r)
            out_k = (k1.unsqueeze(-1) * k2.unsqueeze(-2)).flatten(-2, -1)
            if self.qk_proj is not None:
                out_k = F.linear(out_k, self.qk_proj.to(out_k.dtype))
            if self.k_lora_A is not None:
                lora = F.linear(F.linear(x, self.k_lora_A), self.k_lora_B)
                out_k = out_k + lora.view(B, L, H_kv, hd) * (
                    self.lora_alpha / max(1, self.lora_rank)
                )
            return out_k

        # TPA-FIX-14: checkpointing only when explicitly enabled
        use_ckpt = self.training and self.gradient_checkpointing
        if use_ckpt:
            try:
                from torch.utils.checkpoint import checkpoint

                Q = checkpoint(_compute_q, hidden_states, use_reentrant=False)
                K = checkpoint(_compute_k, hidden_states, use_reentrant=False)
            except Exception:
                Q = _compute_q(hidden_states)
                K = _compute_k(hidden_states)
        else:
            Q = _compute_q(hidden_states)
            K = _compute_k(hidden_states)

        v = self.v_proj(hidden_states).view(B, L, H_kv, hd)

        Q = Q.transpose(1, 2).contiguous()  # (B, H, L, hd)
        K = K.transpose(1, 2).contiguous()
        v = v.transpose(1, 2).contiguous()

        # ---- RoPE ----
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) == 2:
            cos, sin = pos_emb
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # ---- GQA ----
        if H_kv != H:
            K = repeat_kv(K, self.num_kv_groups)
            v = repeat_kv(v, self.num_kv_groups)

        # ---- Mask ----
        is_causal = self.is_decoder and L > 1
        custom_mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask")

        # TPA-FIX-09: preprocess custom_mask WITHOUT losing causality.
        # Normalize custom_mask into a 4D bool mask if it exists.
        if custom_mask is not None:
            if custom_mask.dtype == torch.bool:
                if custom_mask.all():
                    custom_mask = None
                elif custom_mask.dim() == 2:
                    custom_mask = custom_mask.unsqueeze(1).unsqueeze(2)
                elif custom_mask.dim() == 3:
                    custom_mask = custom_mask.unsqueeze(1)
            elif torch.is_floating_point(custom_mask):
                if (custom_mask == 0.0).all():
                    custom_mask = None
                else:
                    # Convert additive float mask to bool
                    if custom_mask.dim() == 2:
                        custom_mask = custom_mask.unsqueeze(1).unsqueeze(2)
                    elif custom_mask.dim() == 3:
                        custom_mask = custom_mask.unsqueeze(1)
                    custom_mask = custom_mask > -1.0  # True = keep

        # TPA-FIX-12: effective head_dim for scaling
        eff_hd = self.effective_qk_dim

        if not output_attentions and head_mask is None:
            # SDPA path
            sdpa_mask = custom_mask
            if is_causal and sdpa_mask is None:
                # Let SDPA handle causality internally
                sdpa_is_causal = True
            else:
                # TPA-FIX-09: build a combined mask = causal AND custom
                if is_causal and sdpa_mask is not None:
                    # Build causal bool mask (B=1, H=1, L, L)
                    causal_bool = (
                        torch.tril(
                            torch.ones(L, L, device=Q.device, dtype=torch.bool),
                        )
                        .unsqueeze(0)
                        .unsqueeze(0)
                    )
                    # Broadcast custom_mask to same shape
                    cm = sdpa_mask
                    if cm.dim() == 2:
                        cm = cm.unsqueeze(1).unsqueeze(2)
                    elif cm.dim() == 3:
                        cm = cm.unsqueeze(1)
                    sdpa_mask = cm & causal_bool
                sdpa_is_causal = False

            try:
                if _HAS_SDPA_KERNEL:
                    with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]):
                        context = F.scaled_dot_product_attention(
                            query=Q,
                            key=K,
                            value=v,
                            attn_mask=sdpa_mask,
                            dropout_p=self.dropout.p if self.training else 0.0,
                            is_causal=sdpa_is_causal,
                        )
                else:
                    context = F.scaled_dot_product_attention(
                        query=Q,
                        key=K,
                        value=v,
                        attn_mask=sdpa_mask,
                        dropout_p=self.dropout.p if self.training else 0.0,
                        is_causal=sdpa_is_causal,
                    )
            except (RuntimeError, TypeError):
                context = F.scaled_dot_product_attention(
                    query=Q,
                    key=K,
                    value=v,
                    attn_mask=sdpa_mask,
                    dropout_p=self.dropout.p if self.training else 0.0,
                    is_causal=sdpa_is_causal,
                )
            attn_weights = None
        else:
            # Manual path
            scale = 1.0 / math.sqrt(eff_hd)
            scores = torch.matmul(Q, K.transpose(-1, -2).contiguous()) * scale

            # TPA-FIX-09: apply causal first, then custom
            if is_causal:
                causal_mask = torch.triu(
                    torch.full((L, L), float("-inf"), device=scores.device, dtype=scores.dtype),
                    diagonal=1,
                )
                scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

            if custom_mask is not None:
                if custom_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~custom_mask, -1e4)
                else:
                    scores = scores + custom_mask.to(scores.dtype)

            attn_weights = F.softmax(scores, dim=-1, dtype=torch.float32).to(v.dtype)
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            attn_weights = self.dropout(attn_weights)

            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)
                attn_weights = attn_weights * head_mask.to(attn_weights.dtype)

            context = torch.matmul(attn_weights, v)

        # ---- Output projection ----
        # TPA-FIX-10: hidden_size matches exactly; no padding hack.
        ctx_flat = context.transpose(1, 2).contiguous().view(B, L, self.q_head_size)
        if self.is_nested_encoder:
            output = ctx_flat
        else:
            output = self.out_proj(ctx_flat)

        # Zero-out padded query positions
        valid_m = get_token_valid_mask(attention_mask, B, L, hidden_states.device)
        if valid_m is not None:
            output = output * valid_m.unsqueeze(-1).to(output.dtype)

        return output, (attn_weights if output_attentions else None)

Every 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 copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    with torch.no_grad():
        w_q = w_k = w_v = w_o = None
        b_q = b_k = b_v = b_o = None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        has_no_out = not any(
            hasattr(old_attn, a)
            for a in ("dense", "out_proj", "o_proj", "c_proj", "out_lin", "proj")
        )
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query") and has_no_out
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight.detach()
                b_o = (
                    old_attn.output.dense.bias.detach()
                    if old_attn.output.dense.bias is not None
                    else None
                )
        else:
            src = old_attn

        # GPT-2
        if hasattr(src, "c_attn"):
            W = src.c_attn.weight.detach()
            B_vec = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
            if W.shape[0] != 3 * self.dim and W.shape[1] == 3 * self.dim:
                W = W.t().contiguous()
            if W.shape[0] == 3 * self.dim:
                w_q, w_k, w_v = W.chunk(3, dim=0)
                if B_vec is not None:
                    b_q, b_k, b_v = B_vec.chunk(3, dim=0)
            if hasattr(src, "c_proj"):
                w_o = src.c_proj.weight.detach()
                if w_o.shape[0] != self.dim and w_o.shape[1] == self.dim:
                    w_o = w_o.t().contiguous()
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # BERT/RoBERTa
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None
            if w_o is None:
                for a in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, a):
                        proj = getattr(src, a)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Qwen / LLaMA
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # DistilBERT
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                if w_q is None:
                    w_q, b_q = adapter.get_weight_bias(src, "q")
                if w_k is None:
                    w_k, b_k = adapter.get_weight_bias(src, "k")
                if w_v is None:
                    w_v, b_v = adapter.get_weight_bias(src, "v")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(src, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        missing = []
        if w_q is None:
            missing.append("Q")
        if w_k is None:
            missing.append("K")
        if w_v is None:
            missing.append("V")
        if w_o is None and not self.is_nested_encoder:
            missing.append("O")
        if missing:
            logger.warning(
                "TPAttention: failed to extract %s from %s",
                missing,
                type(old_attn).__name__,
            )
            return

        # ---- Factorize Q and K ----
        def _factorize_qk(W, proj1, proj2, num_h):
            """TPA-FIX-07/15: try CP; if error > max_cp_error, use SVD."""
            W_f = W.float()
            in_dim = W_f.shape[1]
            r_rows = num_h * self.rank

            cp_ok = False
            if self._use_tensor_product:
                try:
                    A, B = _cp_decompose_qk_fast(
                        W_f,
                        num_h,
                        self.rank,
                        self.head_dim,
                        n_iter=3,
                    )
                    err = _cp_reconstruction_error(
                        W_f,
                        A,
                        B,
                        num_h,
                        self.rank,
                        self.head_dim,
                    )
                    if err <= self.max_cp_error and not (
                        torch.isnan(A).any() or torch.isnan(B).any()
                    ):
                        proj1.weight.data.zero_()
                        proj2.weight.data.zero_()
                        proj1.weight.data[:r_rows, :in_dim].copy_(
                            A[:r_rows, :in_dim].to(proj1.weight.dtype)
                        )
                        proj2.weight.data[:r_rows, :in_dim].copy_(
                            B[:r_rows, :in_dim].to(proj2.weight.dtype)
                        )
                        logger.info("TPAttention: CP OK (err=%.4f) for %d heads", err, num_h)
                        cp_ok = True
                    else:
                        logger.warning(
                            "TPAttention: CP error %.4f > max %.4f; using SVD",
                            err,
                            self.max_cp_error,
                        )
                except Exception as e:
                    logger.warning("TPAttention: CP failed: %s; using SVD", e)

            if not cp_ok:
                # Fallback: use SVD factorized projection, and mark qk_proj
                # to be identity-like (we set proj2 = I to preserve structure).
                # We store SVD factors into proj1 (all rows) and proj2 as identity.
                # This is a slight approximation — the SVD maps W → (U√S, Vt)
                # but the tensor-product forward requires matching shapes.
                # Instead, we use the qk_proj path.
                A, B = _svd_factorize(W_f, self.rank * self.rank)
                # Store full rank factorization in proj1 (as (H*r, in)) and
                # leave proj2 as orthogonal init. The qk_proj path will
                # combine them.
                r_eff = min(A.shape[1], r_rows)
                proj1.weight.data.zero_()
                proj2.weight.data.zero_()
                # Store A directly in proj1 (truncated to r_rows)
                proj1.weight.data[:r_eff, :in_dim].copy_(
                    A[:r_eff, :in_dim].to(proj1.weight.dtype)
                )
                # Store B in proj2 (truncated to r_rows)
                proj2.weight.data[:r_eff, :in_dim].copy_(
                    B[:r_eff, :in_dim].to(proj2.weight.dtype)
                )
                # Set qk_proj to identity-like on the effective dim
                if self.qk_proj is not None:
                    n = min(self.qk_proj.shape)
                    self.qk_proj.data.zero_()
                    for i in range(n):
                        self.qk_proj.data[i, i] = 1.0
                logger.info("TPAttention: SVD fallback applied for %d heads", num_h)

            proj1.weight.requires_grad = False
            proj2.weight.requires_grad = False

        _factorize_qk(w_q, self.q_proj1, self.q_proj2, self.num_heads)
        _factorize_qk(w_k, self.k_proj1, self.k_proj2, self.num_kv_heads)

        # ---- Copy V and O directly ----
        def _copy_layer(target, w, b):
            if w is None:
                return
            if w.shape == target.weight.shape:
                target.weight.copy_(w.to(target.weight.dtype))
            elif w.t().contiguous().shape == target.weight.shape:
                target.weight.copy_(w.t().contiguous().to(target.weight.dtype))
            else:
                m0 = min(target.weight.shape[0], w.shape[0])
                m1 = min(target.weight.shape[1], w.shape[1])
                target.weight.zero_()
                target.weight[:m0, :m1].copy_(w[:m0, :m1].to(target.weight.dtype))
            if b is not None and target.bias is not None:
                mb = min(target.bias.shape[0], b.shape[0])
                target.bias.zero_()
                target.bias[:mb].copy_(b[:mb].to(target.bias.dtype))
            elif target.bias is not None:
                target.bias.zero_()

        _copy_layer(self.v_proj, w_v, b_v)
        if w_o is not None:
            _copy_layer(self.out_proj, w_o, b_o)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif hasattr(old_attn, "c_attn") or any(
            x in old_attn.__class__.__name__.lower()
            for x in ("gpt", "qwen", "llama", "decoder")
        ):
            self.is_decoder = True

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,
*args,
**kwargs) ‑> Callable[..., typing.Any]

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,
    *args,
    **kwargs,
):
    B, L, D = hidden_states.shape
    H, H_kv, r, hd = self.num_heads, self.num_kv_heads, self.rank, self.head_dim

    # ---- Q, K via tensor product ----
    def _compute_q(x):
        q1 = self.q_proj1(x).view(B, L, H, r)
        q2 = self.q_proj2(x).view(B, L, H, r)
        out_q = (q1.unsqueeze(-1) * q2.unsqueeze(-2)).flatten(-2, -1)  # (B, L, H, r*r)
        if self.qk_proj is not None:
            out_q = F.linear(out_q, self.qk_proj.to(out_q.dtype))
        # TPA-FIX-11: add LoRA residual on the (H*hd) space
        if self.q_lora_A is not None:
            lora = F.linear(F.linear(x, self.q_lora_A), self.q_lora_B)
            out_q = out_q + lora.view(B, L, H, hd) * (self.lora_alpha / max(1, self.lora_rank))
        return out_q

    def _compute_k(x):
        k1 = self.k_proj1(x).view(B, L, H_kv, r)
        k2 = self.k_proj2(x).view(B, L, H_kv, r)
        out_k = (k1.unsqueeze(-1) * k2.unsqueeze(-2)).flatten(-2, -1)
        if self.qk_proj is not None:
            out_k = F.linear(out_k, self.qk_proj.to(out_k.dtype))
        if self.k_lora_A is not None:
            lora = F.linear(F.linear(x, self.k_lora_A), self.k_lora_B)
            out_k = out_k + lora.view(B, L, H_kv, hd) * (
                self.lora_alpha / max(1, self.lora_rank)
            )
        return out_k

    # TPA-FIX-14: checkpointing only when explicitly enabled
    use_ckpt = self.training and self.gradient_checkpointing
    if use_ckpt:
        try:
            from torch.utils.checkpoint import checkpoint

            Q = checkpoint(_compute_q, hidden_states, use_reentrant=False)
            K = checkpoint(_compute_k, hidden_states, use_reentrant=False)
        except Exception:
            Q = _compute_q(hidden_states)
            K = _compute_k(hidden_states)
    else:
        Q = _compute_q(hidden_states)
        K = _compute_k(hidden_states)

    v = self.v_proj(hidden_states).view(B, L, H_kv, hd)

    Q = Q.transpose(1, 2).contiguous()  # (B, H, L, hd)
    K = K.transpose(1, 2).contiguous()
    v = v.transpose(1, 2).contiguous()

    # ---- RoPE ----
    pos_emb = kwargs.get("position_embeddings")
    if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) == 2:
        cos, sin = pos_emb
        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

    # ---- GQA ----
    if H_kv != H:
        K = repeat_kv(K, self.num_kv_groups)
        v = repeat_kv(v, self.num_kv_groups)

    # ---- Mask ----
    is_causal = self.is_decoder and L > 1
    custom_mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask")

    # TPA-FIX-09: preprocess custom_mask WITHOUT losing causality.
    # Normalize custom_mask into a 4D bool mask if it exists.
    if custom_mask is not None:
        if custom_mask.dtype == torch.bool:
            if custom_mask.all():
                custom_mask = None
            elif custom_mask.dim() == 2:
                custom_mask = custom_mask.unsqueeze(1).unsqueeze(2)
            elif custom_mask.dim() == 3:
                custom_mask = custom_mask.unsqueeze(1)
        elif torch.is_floating_point(custom_mask):
            if (custom_mask == 0.0).all():
                custom_mask = None
            else:
                # Convert additive float mask to bool
                if custom_mask.dim() == 2:
                    custom_mask = custom_mask.unsqueeze(1).unsqueeze(2)
                elif custom_mask.dim() == 3:
                    custom_mask = custom_mask.unsqueeze(1)
                custom_mask = custom_mask > -1.0  # True = keep

    # TPA-FIX-12: effective head_dim for scaling
    eff_hd = self.effective_qk_dim

    if not output_attentions and head_mask is None:
        # SDPA path
        sdpa_mask = custom_mask
        if is_causal and sdpa_mask is None:
            # Let SDPA handle causality internally
            sdpa_is_causal = True
        else:
            # TPA-FIX-09: build a combined mask = causal AND custom
            if is_causal and sdpa_mask is not None:
                # Build causal bool mask (B=1, H=1, L, L)
                causal_bool = (
                    torch.tril(
                        torch.ones(L, L, device=Q.device, dtype=torch.bool),
                    )
                    .unsqueeze(0)
                    .unsqueeze(0)
                )
                # Broadcast custom_mask to same shape
                cm = sdpa_mask
                if cm.dim() == 2:
                    cm = cm.unsqueeze(1).unsqueeze(2)
                elif cm.dim() == 3:
                    cm = cm.unsqueeze(1)
                sdpa_mask = cm & causal_bool
            sdpa_is_causal = False

        try:
            if _HAS_SDPA_KERNEL:
                with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]):
                    context = F.scaled_dot_product_attention(
                        query=Q,
                        key=K,
                        value=v,
                        attn_mask=sdpa_mask,
                        dropout_p=self.dropout.p if self.training else 0.0,
                        is_causal=sdpa_is_causal,
                    )
            else:
                context = F.scaled_dot_product_attention(
                    query=Q,
                    key=K,
                    value=v,
                    attn_mask=sdpa_mask,
                    dropout_p=self.dropout.p if self.training else 0.0,
                    is_causal=sdpa_is_causal,
                )
        except (RuntimeError, TypeError):
            context = F.scaled_dot_product_attention(
                query=Q,
                key=K,
                value=v,
                attn_mask=sdpa_mask,
                dropout_p=self.dropout.p if self.training else 0.0,
                is_causal=sdpa_is_causal,
            )
        attn_weights = None
    else:
        # Manual path
        scale = 1.0 / math.sqrt(eff_hd)
        scores = torch.matmul(Q, K.transpose(-1, -2).contiguous()) * scale

        # TPA-FIX-09: apply causal first, then custom
        if is_causal:
            causal_mask = torch.triu(
                torch.full((L, L), float("-inf"), device=scores.device, dtype=scores.dtype),
                diagonal=1,
            )
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

        if custom_mask is not None:
            if custom_mask.dtype == torch.bool:
                scores = scores.masked_fill(~custom_mask, -1e4)
            else:
                scores = scores + custom_mask.to(scores.dtype)

        attn_weights = F.softmax(scores, dim=-1, dtype=torch.float32).to(v.dtype)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights)

        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)
            attn_weights = attn_weights * head_mask.to(attn_weights.dtype)

        context = torch.matmul(attn_weights, v)

    # ---- Output projection ----
    # TPA-FIX-10: hidden_size matches exactly; no padding hack.
    ctx_flat = context.transpose(1, 2).contiguous().view(B, L, self.q_head_size)
    if self.is_nested_encoder:
        output = ctx_flat
    else:
        output = self.out_proj(ctx_flat)

    # Zero-out padded query positions
    valid_m = get_token_valid_mask(attention_mask, B, L, hidden_states.device)
    if valid_m is not None:
        output = output * valid_m.unsqueeze(-1).to(output.dtype)

    return output, (attn_weights if output_attentions else None)

Must return (output, attention_weights_or_None).

class TPAttention (config: Any | None = None,
rank: int = 8,
pretrained_attn: nn.Module | None = None,
lora_rank: int = 0,
lora_alpha: float = 1.0,
is_decoder: bool | None = None,
max_cp_error: float = 0.35,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("tpa", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("tpa_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("tp_attention", supports_causal=True, supports_gqa=True)
class TPAttention(BaseAttention):
    supports_gqa: bool = True
    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 8,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 0,
        lora_alpha: float = 1.0,
        is_decoder: bool | None = None,
        max_cp_error: float = 0.35,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        src = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")
        self.max_cp_error = float(max_cp_error)

        # ---- hidden size ----
        dim = None
        if config is not None and not isinstance(config, int):
            dim = _cfg_get(
                config,
                "hidden_size",
                _cfg_get(config, "dim", _cfg_get(config, "d_model", None)),
            )
        if dim is None and src is not None:
            dim = getattr(src, "hidden_size", getattr(src, "embed_dim", None))
            if dim is None:
                for attr in ("q_proj", "q_lin", "query"):
                    m = getattr(src, attr, None)
                    if m is not None and hasattr(m, "in_features"):
                        dim = m.in_features
                        break
            if dim is None and hasattr(src, "c_attn") and hasattr(src.c_attn, "weight"):
                w = src.c_attn.weight
                dim = w.shape[0] if w.shape[0] != 3 * w.shape[1] else w.shape[1]
        if dim is None:
            dim = kwargs.get("hidden_size", 768)

        # ---- heads ----
        num_heads = None
        if config is not None and not isinstance(config, int):
            num_heads = _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "n_heads", _cfg_get(config, "num_heads", None)),
            )
        if num_heads is None and src is not None:
            num_heads = getattr(
                src,
                "num_heads",
                getattr(src, "num_attention_heads", getattr(src, "n_head", None)),
            )
        if num_heads is None:
            num_heads = 12

        self.dim = int(dim)
        self.num_heads = max(1, int(num_heads))

        # ---- head_dim ----
        cfg_hd = _cfg_get(config, "head_dim", None) if config is not None else None
        kw_hd = kwargs.get("head_dim")
        if cfg_hd is not None:
            self.head_dim = int(cfg_hd)
        elif kw_hd is not None:
            self.head_dim = int(kw_hd)
        elif src is not None and getattr(src, "head_dim", None) is not None:
            self.head_dim = int(src.head_dim)
        else:
            self.head_dim = self.dim // self.num_heads

        # ---- GQA ----
        cfg_kv = (
            _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
            if config is not None
            else None
        )
        if cfg_kv is None and src is not None:
            cfg_kv = getattr(src, "num_key_value_heads", getattr(src, "num_kv_heads", None))
        self.num_kv_heads = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_kv_groups = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.q_head_size = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim

        # ---- rank (head-rank for tensor product) ----
        raw_r = kwargs.pop("rank", _cfg_get(config, "rank", rank))
        # TPA-FIX-02: rank must satisfy r*r == head_dim for tensor product.
        # If not, we use the qk_proj path (proper projection).
        self.rank = int(raw_r)
        self._use_tensor_product = self.rank * self.rank == self.head_dim
        # Effective qk dimension after tensor-product or qk_proj
        self.effective_qk_dim = self.rank * self.rank if self._use_tensor_product else self.head_dim

        self.lora_rank = int(kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank)))
        self.lora_alpha = float(
            kwargs.get("lora_alpha", _cfg_get(config, "lora_alpha", lora_alpha))
        )

        # ---- decoder ----
        _mt = str(_cfg_get(config, "model_type", "") or "").lower() if config is not None else ""
        _dec_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        elif config is not None and (
            _cfg_get(config, "is_decoder", False)
            or _cfg_get(config, "is_causal", False)
            or any(_mt.startswith(p) for p in _dec_prefixes)
        ):
            self.is_decoder = True
        else:
            self.is_decoder = bool(kwargs.get("is_decoder", False))

        self.is_nested_encoder = False
        # TPA-FIX-14: gradient checkpointing only when explicitly requested.
        self.gradient_checkpointing = bool(kwargs.get("gradient_checkpointing", False))

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and src is not None:
            p = next(src.parameters(), None)
            if p is not None:
                target_dtype = p.dtype

        # ---- projections ----
        # TPA-FIX-11: LoRA is applied to a separate low-rank residual, not to
        # the tensor-product factors. This keeps the semantics well-defined.
        self.q_proj1 = LoRALinear(
            self.dim,
            self.num_heads * self.rank,
            bias=False,
            lora_rank=0,
            lora_alpha=1.0,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.q_proj2 = LoRALinear(
            self.dim,
            self.num_heads * self.rank,
            bias=False,
            lora_rank=0,
            lora_alpha=1.0,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.k_proj1 = LoRALinear(
            self.dim,
            self.num_kv_heads * self.rank,
            bias=False,
            lora_rank=0,
            lora_alpha=1.0,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.k_proj2 = LoRALinear(
            self.dim,
            self.num_kv_heads * self.rank,
            bias=False,
            lora_rank=0,
            lora_alpha=1.0,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.out_proj = LoRALinear(
            self.q_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )

        # TPA-FIX-08/15: when not using tensor-product, use a proper projection
        # from (rank*rank) → head_dim. Otherwise the tensor product already
        # produces head_dim directly.
        if not self._use_tensor_product:
            n_in = self.rank * self.rank
            n_out = self.head_dim
            self.qk_proj = nn.Parameter(torch.zeros(n_out, n_in), requires_grad=False)
        else:
            self.register_parameter("qk_proj", None)

        # TPA-FIX-11: dedicated LoRA on the effective Q/K projections.
        # We store them as LoRALinear to align with the framework convention.
        # They apply to the concatenated (H*hd) space, not the factorized dims.
        if self.lora_rank > 0:
            self.q_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.dim, dtype=target_dtype))
            self.q_lora_B = nn.Parameter(
                torch.zeros(self.q_head_size, self.lora_rank, dtype=target_dtype)
            )
            self.k_lora_A = nn.Parameter(torch.empty(self.lora_rank, self.dim, dtype=target_dtype))
            self.k_lora_B = nn.Parameter(
                torch.zeros(self.kv_head_size, self.lora_rank, dtype=target_dtype)
            )
            nn.init.kaiming_uniform_(self.q_lora_A, a=math.sqrt(5))
            nn.init.kaiming_uniform_(self.k_lora_A, a=math.sqrt(5))
        else:
            self.register_parameter("q_lora_A", None)
            self.register_parameter("q_lora_B", None)
            self.register_parameter("k_lora_A", None)
            self.register_parameter("k_lora_B", None)

        dropout_p = (
            getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1))
            if config
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

        if target_dtype is not None:
            for proj in (
                self.q_proj1,
                self.q_proj2,
                self.k_proj1,
                self.k_proj2,
                self.v_proj,
                self.out_proj,
            ):
                proj.to(target_dtype)

    # ------------------------------------------------------------------
    # Weight transfer
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        with torch.no_grad():
            w_q = w_k = w_v = w_o = None
            b_q = b_k = b_v = b_o = None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            has_no_out = not any(
                hasattr(old_attn, a)
                for a in ("dense", "out_proj", "o_proj", "c_proj", "out_lin", "proj")
            )
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query") and has_no_out
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight.detach()
                    b_o = (
                        old_attn.output.dense.bias.detach()
                        if old_attn.output.dense.bias is not None
                        else None
                    )
            else:
                src = old_attn

            # GPT-2
            if hasattr(src, "c_attn"):
                W = src.c_attn.weight.detach()
                B_vec = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                if W.shape[0] != 3 * self.dim and W.shape[1] == 3 * self.dim:
                    W = W.t().contiguous()
                if W.shape[0] == 3 * self.dim:
                    w_q, w_k, w_v = W.chunk(3, dim=0)
                    if B_vec is not None:
                        b_q, b_k, b_v = B_vec.chunk(3, dim=0)
                if hasattr(src, "c_proj"):
                    w_o = src.c_proj.weight.detach()
                    if w_o.shape[0] != self.dim and w_o.shape[1] == self.dim:
                        w_o = w_o.t().contiguous()
                    b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # BERT/RoBERTa
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None
                if w_o is None:
                    for a in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, a):
                            proj = getattr(src, a)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Qwen / LLaMA
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # DistilBERT
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    if w_q is None:
                        w_q, b_q = adapter.get_weight_bias(src, "q")
                    if w_k is None:
                        w_k, b_k = adapter.get_weight_bias(src, "k")
                    if w_v is None:
                        w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            missing = []
            if w_q is None:
                missing.append("Q")
            if w_k is None:
                missing.append("K")
            if w_v is None:
                missing.append("V")
            if w_o is None and not self.is_nested_encoder:
                missing.append("O")
            if missing:
                logger.warning(
                    "TPAttention: failed to extract %s from %s",
                    missing,
                    type(old_attn).__name__,
                )
                return

            # ---- Factorize Q and K ----
            def _factorize_qk(W, proj1, proj2, num_h):
                """TPA-FIX-07/15: try CP; if error > max_cp_error, use SVD."""
                W_f = W.float()
                in_dim = W_f.shape[1]
                r_rows = num_h * self.rank

                cp_ok = False
                if self._use_tensor_product:
                    try:
                        A, B = _cp_decompose_qk_fast(
                            W_f,
                            num_h,
                            self.rank,
                            self.head_dim,
                            n_iter=3,
                        )
                        err = _cp_reconstruction_error(
                            W_f,
                            A,
                            B,
                            num_h,
                            self.rank,
                            self.head_dim,
                        )
                        if err <= self.max_cp_error and not (
                            torch.isnan(A).any() or torch.isnan(B).any()
                        ):
                            proj1.weight.data.zero_()
                            proj2.weight.data.zero_()
                            proj1.weight.data[:r_rows, :in_dim].copy_(
                                A[:r_rows, :in_dim].to(proj1.weight.dtype)
                            )
                            proj2.weight.data[:r_rows, :in_dim].copy_(
                                B[:r_rows, :in_dim].to(proj2.weight.dtype)
                            )
                            logger.info("TPAttention: CP OK (err=%.4f) for %d heads", err, num_h)
                            cp_ok = True
                        else:
                            logger.warning(
                                "TPAttention: CP error %.4f > max %.4f; using SVD",
                                err,
                                self.max_cp_error,
                            )
                    except Exception as e:
                        logger.warning("TPAttention: CP failed: %s; using SVD", e)

                if not cp_ok:
                    # Fallback: use SVD factorized projection, and mark qk_proj
                    # to be identity-like (we set proj2 = I to preserve structure).
                    # We store SVD factors into proj1 (all rows) and proj2 as identity.
                    # This is a slight approximation — the SVD maps W → (U√S, Vt)
                    # but the tensor-product forward requires matching shapes.
                    # Instead, we use the qk_proj path.
                    A, B = _svd_factorize(W_f, self.rank * self.rank)
                    # Store full rank factorization in proj1 (as (H*r, in)) and
                    # leave proj2 as orthogonal init. The qk_proj path will
                    # combine them.
                    r_eff = min(A.shape[1], r_rows)
                    proj1.weight.data.zero_()
                    proj2.weight.data.zero_()
                    # Store A directly in proj1 (truncated to r_rows)
                    proj1.weight.data[:r_eff, :in_dim].copy_(
                        A[:r_eff, :in_dim].to(proj1.weight.dtype)
                    )
                    # Store B in proj2 (truncated to r_rows)
                    proj2.weight.data[:r_eff, :in_dim].copy_(
                        B[:r_eff, :in_dim].to(proj2.weight.dtype)
                    )
                    # Set qk_proj to identity-like on the effective dim
                    if self.qk_proj is not None:
                        n = min(self.qk_proj.shape)
                        self.qk_proj.data.zero_()
                        for i in range(n):
                            self.qk_proj.data[i, i] = 1.0
                    logger.info("TPAttention: SVD fallback applied for %d heads", num_h)

                proj1.weight.requires_grad = False
                proj2.weight.requires_grad = False

            _factorize_qk(w_q, self.q_proj1, self.q_proj2, self.num_heads)
            _factorize_qk(w_k, self.k_proj1, self.k_proj2, self.num_kv_heads)

            # ---- Copy V and O directly ----
            def _copy_layer(target, w, b):
                if w is None:
                    return
                if w.shape == target.weight.shape:
                    target.weight.copy_(w.to(target.weight.dtype))
                elif w.t().contiguous().shape == target.weight.shape:
                    target.weight.copy_(w.t().contiguous().to(target.weight.dtype))
                else:
                    m0 = min(target.weight.shape[0], w.shape[0])
                    m1 = min(target.weight.shape[1], w.shape[1])
                    target.weight.zero_()
                    target.weight[:m0, :m1].copy_(w[:m0, :m1].to(target.weight.dtype))
                if b is not None and target.bias is not None:
                    mb = min(target.bias.shape[0], b.shape[0])
                    target.bias.zero_()
                    target.bias[:mb].copy_(b[:mb].to(target.bias.dtype))
                elif target.bias is not None:
                    target.bias.zero_()

            _copy_layer(self.v_proj, w_v, b_v)
            if w_o is not None:
                _copy_layer(self.out_proj, w_o, b_o)

            if hasattr(old_attn, "is_decoder"):
                self.is_decoder = bool(old_attn.is_decoder)
            elif hasattr(old_attn, "c_attn") or any(
                x in old_attn.__class__.__name__.lower()
                for x in ("gpt", "qwen", "llama", "decoder")
            ):
                self.is_decoder = True

    # ------------------------------------------------------------------
    # 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,
        *args,
        **kwargs,
    ):
        B, L, D = hidden_states.shape
        H, H_kv, r, hd = self.num_heads, self.num_kv_heads, self.rank, self.head_dim

        # ---- Q, K via tensor product ----
        def _compute_q(x):
            q1 = self.q_proj1(x).view(B, L, H, r)
            q2 = self.q_proj2(x).view(B, L, H, r)
            out_q = (q1.unsqueeze(-1) * q2.unsqueeze(-2)).flatten(-2, -1)  # (B, L, H, r*r)
            if self.qk_proj is not None:
                out_q = F.linear(out_q, self.qk_proj.to(out_q.dtype))
            # TPA-FIX-11: add LoRA residual on the (H*hd) space
            if self.q_lora_A is not None:
                lora = F.linear(F.linear(x, self.q_lora_A), self.q_lora_B)
                out_q = out_q + lora.view(B, L, H, hd) * (self.lora_alpha / max(1, self.lora_rank))
            return out_q

        def _compute_k(x):
            k1 = self.k_proj1(x).view(B, L, H_kv, r)
            k2 = self.k_proj2(x).view(B, L, H_kv, r)
            out_k = (k1.unsqueeze(-1) * k2.unsqueeze(-2)).flatten(-2, -1)
            if self.qk_proj is not None:
                out_k = F.linear(out_k, self.qk_proj.to(out_k.dtype))
            if self.k_lora_A is not None:
                lora = F.linear(F.linear(x, self.k_lora_A), self.k_lora_B)
                out_k = out_k + lora.view(B, L, H_kv, hd) * (
                    self.lora_alpha / max(1, self.lora_rank)
                )
            return out_k

        # TPA-FIX-14: checkpointing only when explicitly enabled
        use_ckpt = self.training and self.gradient_checkpointing
        if use_ckpt:
            try:
                from torch.utils.checkpoint import checkpoint

                Q = checkpoint(_compute_q, hidden_states, use_reentrant=False)
                K = checkpoint(_compute_k, hidden_states, use_reentrant=False)
            except Exception:
                Q = _compute_q(hidden_states)
                K = _compute_k(hidden_states)
        else:
            Q = _compute_q(hidden_states)
            K = _compute_k(hidden_states)

        v = self.v_proj(hidden_states).view(B, L, H_kv, hd)

        Q = Q.transpose(1, 2).contiguous()  # (B, H, L, hd)
        K = K.transpose(1, 2).contiguous()
        v = v.transpose(1, 2).contiguous()

        # ---- RoPE ----
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) == 2:
            cos, sin = pos_emb
            Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # ---- GQA ----
        if H_kv != H:
            K = repeat_kv(K, self.num_kv_groups)
            v = repeat_kv(v, self.num_kv_groups)

        # ---- Mask ----
        is_causal = self.is_decoder and L > 1
        custom_mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask")

        # TPA-FIX-09: preprocess custom_mask WITHOUT losing causality.
        # Normalize custom_mask into a 4D bool mask if it exists.
        if custom_mask is not None:
            if custom_mask.dtype == torch.bool:
                if custom_mask.all():
                    custom_mask = None
                elif custom_mask.dim() == 2:
                    custom_mask = custom_mask.unsqueeze(1).unsqueeze(2)
                elif custom_mask.dim() == 3:
                    custom_mask = custom_mask.unsqueeze(1)
            elif torch.is_floating_point(custom_mask):
                if (custom_mask == 0.0).all():
                    custom_mask = None
                else:
                    # Convert additive float mask to bool
                    if custom_mask.dim() == 2:
                        custom_mask = custom_mask.unsqueeze(1).unsqueeze(2)
                    elif custom_mask.dim() == 3:
                        custom_mask = custom_mask.unsqueeze(1)
                    custom_mask = custom_mask > -1.0  # True = keep

        # TPA-FIX-12: effective head_dim for scaling
        eff_hd = self.effective_qk_dim

        if not output_attentions and head_mask is None:
            # SDPA path
            sdpa_mask = custom_mask
            if is_causal and sdpa_mask is None:
                # Let SDPA handle causality internally
                sdpa_is_causal = True
            else:
                # TPA-FIX-09: build a combined mask = causal AND custom
                if is_causal and sdpa_mask is not None:
                    # Build causal bool mask (B=1, H=1, L, L)
                    causal_bool = (
                        torch.tril(
                            torch.ones(L, L, device=Q.device, dtype=torch.bool),
                        )
                        .unsqueeze(0)
                        .unsqueeze(0)
                    )
                    # Broadcast custom_mask to same shape
                    cm = sdpa_mask
                    if cm.dim() == 2:
                        cm = cm.unsqueeze(1).unsqueeze(2)
                    elif cm.dim() == 3:
                        cm = cm.unsqueeze(1)
                    sdpa_mask = cm & causal_bool
                sdpa_is_causal = False

            try:
                if _HAS_SDPA_KERNEL:
                    with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]):
                        context = F.scaled_dot_product_attention(
                            query=Q,
                            key=K,
                            value=v,
                            attn_mask=sdpa_mask,
                            dropout_p=self.dropout.p if self.training else 0.0,
                            is_causal=sdpa_is_causal,
                        )
                else:
                    context = F.scaled_dot_product_attention(
                        query=Q,
                        key=K,
                        value=v,
                        attn_mask=sdpa_mask,
                        dropout_p=self.dropout.p if self.training else 0.0,
                        is_causal=sdpa_is_causal,
                    )
            except (RuntimeError, TypeError):
                context = F.scaled_dot_product_attention(
                    query=Q,
                    key=K,
                    value=v,
                    attn_mask=sdpa_mask,
                    dropout_p=self.dropout.p if self.training else 0.0,
                    is_causal=sdpa_is_causal,
                )
            attn_weights = None
        else:
            # Manual path
            scale = 1.0 / math.sqrt(eff_hd)
            scores = torch.matmul(Q, K.transpose(-1, -2).contiguous()) * scale

            # TPA-FIX-09: apply causal first, then custom
            if is_causal:
                causal_mask = torch.triu(
                    torch.full((L, L), float("-inf"), device=scores.device, dtype=scores.dtype),
                    diagonal=1,
                )
                scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

            if custom_mask is not None:
                if custom_mask.dtype == torch.bool:
                    scores = scores.masked_fill(~custom_mask, -1e4)
                else:
                    scores = scores + custom_mask.to(scores.dtype)

            attn_weights = F.softmax(scores, dim=-1, dtype=torch.float32).to(v.dtype)
            attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
            attn_weights = self.dropout(attn_weights)

            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)
                attn_weights = attn_weights * head_mask.to(attn_weights.dtype)

            context = torch.matmul(attn_weights, v)

        # ---- Output projection ----
        # TPA-FIX-10: hidden_size matches exactly; no padding hack.
        ctx_flat = context.transpose(1, 2).contiguous().view(B, L, self.q_head_size)
        if self.is_nested_encoder:
            output = ctx_flat
        else:
            output = self.out_proj(ctx_flat)

        # Zero-out padded query positions
        valid_m = get_token_valid_mask(attention_mask, B, L, hidden_states.device)
        if valid_m is not None:
            output = output * valid_m.unsqueeze(-1).to(output.dtype)

        return output, (attn_weights if output_attentions else None)

Every 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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class TokenCompressedAttention (config=None,
rank=64,
compression_ratio=0.5,
pretrained_attn=None,
lora_rank=None,
lora_alpha=None,
comp_config=None,
is_decoder=None,
**kwargs)

Expand source code

@ATTENTION_REGISTRY.register("token_compressed")
@ATTENTION_REGISTRY.register("token_compressed_attention")
class TokenCompressedAttention(BaseAttention):
    supports_gqa = True
    supports_causal = True
    is_custom_attention = True

    def __init__(
        self,
        config=None,
        rank=64,
        compression_ratio=0.5,
        pretrained_attn=None,
        lora_rank=None,
        lora_alpha=None,
        comp_config=None,
        is_decoder=None,
        **kwargs,
    ):
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        src = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")

        # ---- hidden_size ----
        dim = None
        if config is not None and not isinstance(config, int):
            dim = _cfg_get(
                config,
                "hidden_size",
                _cfg_get(config, "dim", _cfg_get(config, "d_model", None)),
            )
        if dim is None and comp_config is not None:
            dim = _cfg_get(comp_config, "hidden_size", _cfg_get(comp_config, "dim", None))
        if dim is None and src is not None:
            dim = getattr(src, "hidden_size", getattr(src, "embed_dim", None))
            if dim is None:
                for attr in ("q_proj", "q_lin", "query"):
                    m = getattr(src, attr, None)
                    if m is not None and hasattr(m, "in_features"):
                        dim = m.in_features
                        break
            if dim is None and hasattr(src, "c_attn"):
                w = src.c_attn.weight
                dim = w.shape[0] if w.shape[0] < w.shape[1] else w.shape[1]
        if dim is None:
            dim = kwargs.get("hidden_size", 768)
        self.dim = int(dim)
        self.hidden_size = self.dim

        # ---- num_heads ----
        num_heads = kwargs.get("num_heads")
        if num_heads is None and config is not None and not isinstance(config, int):
            num_heads = _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "n_heads", _cfg_get(config, "num_heads", None)),
            )
        if num_heads is None and comp_config is not None:
            num_heads = _cfg_get(
                comp_config,
                "num_attention_heads",
                _cfg_get(comp_config, "n_heads", None),
            )
        if num_heads is None and src is not None:
            num_heads = getattr(
                src,
                "num_heads",
                getattr(src, "num_attention_heads", getattr(src, "n_head", None)),
            )
        if num_heads is None:
            num_heads = max(1, self.dim // 64)
        self.num_heads = max(1, int(num_heads))

        # ---- head_dim ----
        head_dim = None
        if config is not None and not isinstance(config, int):
            head_dim = _cfg_get(config, "head_dim", None)
        if head_dim is None:
            head_dim = kwargs.get("head_dim")
        if head_dim is None:
            head_dim = self.dim // self.num_heads
        self.head_dim = int(head_dim)

        # ---- num_kv_heads ----
        cfg_kv = None
        if config is not None and not isinstance(config, int):
            cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is None and comp_config is not None:
            cfg_kv = _cfg_get(
                comp_config,
                "num_key_value_heads",
                _cfg_get(comp_config, "num_kv_heads", None),
            )
        if cfg_kv is None and src is not None:
            cfg_kv = getattr(src, "num_key_value_heads", getattr(src, "num_kv_heads", None))
        self.num_kv_heads = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.num_key_value_groups = max(1, self.num_heads // self.num_kv_heads)

        self.q_head_size = self.num_heads * self.head_dim
        self.kv_head_size = self.num_kv_heads * self.head_dim
        self.all_head_size = self.q_head_size
        self.all_kv_head_size = self.kv_head_size

        # ---- compression ratio ----
        self.compression_ratio = float(
            kwargs.get(
                "compression_ratio",
                _cfg_get(comp_config, "compression_ratio", compression_ratio),
            )
        )
        self.compression_ratio = max(0.1, min(1.0, self.compression_ratio))

        # ---- LoRA config ----
        rl = kwargs.get("lora_rank")
        if rl is None and comp_config is not None:
            rl = _cfg_get(comp_config, "lora_rank")
        if rl is None and config is not None and not isinstance(config, int):
            rl = _cfg_get(config, "lora_rank")
        if rl is None:
            rl = lora_rank if lora_rank is not None else 16
        self.lora_rank = int(rl)

        ra = kwargs.get("lora_alpha")
        if ra is None and comp_config is not None:
            ra = _cfg_get(comp_config, "lora_alpha")
        if ra is None and config is not None and not isinstance(config, int):
            ra = _cfg_get(config, "lora_alpha")
        if ra is None:
            ra = lora_alpha if lora_alpha is not None else 32.0
        self.lora_alpha = float(ra)

        # ---- decoder detection ----
        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        elif kwargs.get("is_decoder") is not None:
            self.is_decoder = bool(kwargs["is_decoder"])
        elif src is not None and getattr(src, "is_decoder", False):
            self.is_decoder = True
        elif config is not None and not isinstance(config, int):
            mt = str(_cfg_get(config, "model_type", "")).lower()
            self.is_decoder = bool(_cfg_get(config, "is_decoder", False)) or any(
                mt.startswith(p)
                for p in (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "gemma",
                    "phi",
                    "opt",
                    "bloom",
                    "falcon",
                    "olmo",
                    "mpt",
                    "stablelm",
                )
            )
        else:
            self.is_decoder = False
        self.is_causal = self.is_decoder
        self.is_nested_encoder = False
        self.rotary_emb = getattr(src, "rotary_emb", None) if src is not None else None

        # ---- device/dtype from source ----
        proj_device = None
        proj_dtype = None
        if src is not None:
            for attr in ("q_lin", "q_proj", "query", "c_attn"):
                mod = getattr(src, attr, None)
                if mod is not None and hasattr(mod, "weight"):
                    proj_device = mod.weight.device
                    proj_dtype = mod.weight.dtype
                    break

        # ---- projections ----
        self.q_lin = LoRALinear(
            self.dim,
            self.q_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            device=proj_device,
            dtype=proj_dtype,
        )
        self.k_lin = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            device=proj_device,
            dtype=proj_dtype,
        )
        self.v_lin = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            device=proj_device,
            dtype=proj_dtype,
        )
        self.out_lin = LoRALinear(
            self.q_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            device=proj_device,
            dtype=proj_dtype,
        )

        self.q_proj = self.q_lin
        self.k_proj = self.k_lin
        self.v_proj = self.v_lin
        self.out_proj = self.out_lin
        self.query = self.q_lin
        self.key = self.k_lin
        self.value = self.v_lin

        # ---- score_net (TC-24: zero-init output) ----
        self.score_net = nn.Sequential(
            nn.Linear(self.dim, 64, device=proj_device, dtype=proj_dtype),
            nn.GELU(),
            nn.Linear(64, 1, device=proj_device, dtype=proj_dtype),
        )
        with torch.no_grad():
            # Zero last layer so score_net starts contributing zero.
            nn.init.zeros_(self.score_net[2].weight)
            nn.init.zeros_(self.score_net[2].bias)

        dropout_p = (
            _cfg_get(config, "dropout", _cfg_get(config, "attention_probs_dropout_prob", 0.1))
            if config is not None and not isinstance(config, int)
            else 0.1
        )
        self.dropout = nn.Dropout(dropout_p)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    # ---- aliases ----
    @property
    def dense(self):
        return self.out_lin

    @property
    def c_proj(self):
        return self.out_lin

    @property
    def o_proj(self):
        return self.out_lin

    @property
    def q_lin_alias(self):
        return self.q_lin

    @property
    def out_proj_alias(self):
        return self.out_lin

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn):
        with torch.no_grad():
            if getattr(old_attn, "rotary_emb", None) is not None:
                self.rotary_emb = old_attn.rotary_emb

            w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
                and not hasattr(old_attn, "c_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    dense = getattr(old_attn.output.dense, "base_layer", old_attn.output.dense)
                    w_o = dense.weight.detach()
                    b_o = dense.bias.detach() if dense.bias is not None else None
            else:
                src = old_attn

            def _unwrap(layer):
                return getattr(layer, "base_layer", layer)

            # GPT-2 fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn = _unwrap(src.c_attn)
                c_proj = _unwrap(src.c_proj)
                cw = c_attn.weight.detach()
                cb = c_attn.bias.detach() if c_attn.bias is not None else None
                if cw.dim() == 2 and cw.shape[0] == self.dim and cw.shape[1] != self.dim:
                    cw = cw.t().contiguous()
                qd, kd, vd = self.q_head_size, self.kv_head_size, self.kv_head_size
                w_q = cw[:qd]
                w_k = cw[qd : qd + kd]
                w_v = cw[qd + kd : qd + kd + vd]
                if cb is not None:
                    b_q = cb[:qd]
                    b_k = cb[qd : qd + kd]
                    b_v = cb[qd + kd : qd + kd + vd]
                ow = c_proj.weight.detach()
                if ow.dim() == 2 and ow.shape[0] == self.dim and ow.shape[1] != self.dim:
                    ow = ow.t().contiguous()
                w_o = ow
                b_o = c_proj.bias.detach() if c_proj.bias is not None else None

            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                q_l, k_l, v_l = _unwrap(src.query), _unwrap(src.key), _unwrap(src.value)
                w_q = q_l.weight.detach()
                b_q = q_l.bias.detach() if q_l.bias is not None else None
                w_k = k_l.weight.detach()
                b_k = k_l.bias.detach() if k_l.bias is not None else None
                w_v = v_l.weight.detach()
                b_v = v_l.bias.detach() if v_l.bias is not None else None
                if w_o is None:
                    for a in ("dense", "out_proj", "o_proj", "c_proj"):
                        if hasattr(src, a):
                            p = _unwrap(getattr(src, a))
                            w_o = p.weight.detach()
                            b_o = p.bias.detach() if p.bias is not None else None
                            break

            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                q_l, k_l, v_l = (
                    _unwrap(src.q_proj),
                    _unwrap(src.k_proj),
                    _unwrap(src.v_proj),
                )
                w_q = q_l.weight.detach()
                b_q = q_l.bias.detach() if q_l.bias is not None else None
                w_k = k_l.weight.detach()
                b_k = k_l.bias.detach() if k_l.bias is not None else None
                w_v = v_l.weight.detach()
                b_v = v_l.bias.detach() if v_l.bias is not None else None
                if hasattr(src, "o_proj"):
                    o_l = _unwrap(src.o_proj)
                    w_o = o_l.weight.detach()
                    b_o = o_l.bias.detach() if o_l.bias is not None else None

            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                q_l, k_l, v_l = (
                    _unwrap(src.q_lin),
                    _unwrap(src.k_lin),
                    _unwrap(src.v_lin),
                )
                w_q = q_l.weight.detach()
                b_q = q_l.bias.detach() if q_l.bias is not None else None
                w_k = k_l.weight.detach()
                b_k = k_l.bias.detach() if k_l.bias is not None else None
                w_v = v_l.weight.detach()
                b_v = v_l.bias.detach() if v_l.bias is not None else None
                if hasattr(src, "out_lin"):
                    o_l = _unwrap(src.out_lin)
                    w_o = o_l.weight.detach()
                    b_o = o_l.bias.detach() if o_l.bias is not None else None

            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    if w_q is None:
                        w_q, b_q = adapter.get_weight_bias(src, "q")
                    if w_k is None:
                        w_k, b_k = adapter.get_weight_bias(src, "k")
                    if w_v is None:
                        w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            missing = []
            if w_q is None:
                missing.append("Q")
            if w_k is None:
                missing.append("K")
            if w_v is None:
                missing.append("V")
            if w_o is None and not self.is_nested_encoder:
                missing.append("O")
            if missing:
                raise RuntimeError(
                    f"TokenCompressedAttention.copy_weights: missing {missing} from {type(old_attn).__name__}"
                )

            def _copy(target_layer, w, b):
                if w is not None:
                    if w.shape == target_layer.weight.shape:
                        target_layer.weight.data.copy_(w.to(target_layer.weight.dtype))
                    elif w.t().contiguous().shape == target_layer.weight.shape:
                        target_layer.weight.data.copy_(
                            w.t().contiguous().to(target_layer.weight.dtype)
                        )
                    else:
                        m0 = min(target_layer.weight.shape[0], w.shape[0])
                        m1 = min(target_layer.weight.shape[1], w.shape[1])
                        target_layer.weight.data.zero_()
                        target_layer.weight.data[:m0, :m1].copy_(
                            w[:m0, :m1].to(target_layer.weight.dtype)
                        )
                if b is not None and target_layer.bias is not None:
                    if b.shape == target_layer.bias.shape:
                        target_layer.bias.data.copy_(b.to(target_layer.bias.dtype))
                    else:
                        mb = min(target_layer.bias.shape[0], b.shape[0])
                        target_layer.bias.data.zero_()
                        target_layer.bias.data[:mb].copy_(b[:mb].to(target_layer.bias.dtype))
                elif target_layer.bias is not None:
                    target_layer.bias.data.zero_()
                target_layer.weight.requires_grad = False
                if target_layer.bias is not None:
                    target_layer.bias.requires_grad = False

            _copy(self.q_lin, w_q, b_q)
            _copy(self.k_lin, w_k, b_k)
            _copy(self.v_lin, w_v, b_v)
            if w_o is not None:
                _copy(self.out_lin, w_o, b_o)

    # ------------------------------------------------------------------
    # Token selection
    # ------------------------------------------------------------------
    def _select_compressed_tokens(
        self,
        x,
        k_tensor,
        v_tensor,
        valid_mask_2d,
    ):
        B, L, D = x.shape
        H_kv, Dh = self.num_kv_heads, self.head_dim

        if L <= 2:
            top_indices = torch.arange(L, device=x.device).unsqueeze(0).expand(B, L)
            comp_mask = valid_mask_2d if valid_mask_2d is not None else None
            return k_tensor, v_tensor, top_indices, comp_mask

        x_f = x.float()
        score_dev = next(self.score_net.parameters()).device
        score_dtype = next(self.score_net.parameters()).dtype
        energy = torch.norm(x_f, p=2, dim=-1)
        learned = (
            self.score_net(x_f.to(score_dev, score_dtype)).squeeze(-1).to(x_f.device, x_f.dtype)
        )
        total_importance = energy + learned

        k_keep = max(2, int(math.ceil(L * self.compression_ratio)))
        k_keep = min(k_keep, L)

        # Mask padding
        if valid_mask_2d is not None:
            total_importance = total_importance.masked_fill(~valid_mask_2d, -1e4)

        # TC-25: if all masked, force keep index 0
        all_masked = (total_importance <= -1e4).all(dim=-1)  # (B,)
        if all_masked.any():
            total_importance = total_importance.clone()
            total_importance[all_masked, 0] = 0.0

        # TC-22: always keep CLS for decoders
        if self.is_decoder:
            cls_boost = torch.full((B, 1), 1e4, device=x.device, dtype=total_importance.dtype)
            total_importance = torch.cat([cls_boost, total_importance[:, 1:]], dim=-1)

        _, top_indices = torch.topk(total_importance, k_keep, dim=-1)
        top_indices, _ = torch.sort(top_indices, dim=-1)

        top_idx_exp = top_indices.unsqueeze(1).unsqueeze(-1).expand(-1, H_kv, -1, Dh)
        k_comp = torch.gather(k_tensor, 2, top_idx_exp)
        v_comp = torch.gather(v_tensor, 2, top_idx_exp)

        comp_mask = None
        if valid_mask_2d is not None:
            comp_mask = torch.gather(valid_mask_2d, 1, top_indices)

        return k_comp, v_comp, top_indices, comp_mask

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_value=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
        position_embeddings=None,
        layer_idx=0,
        *args,
        **kwargs,
    ):
        if encoder_hidden_states is not None:
            raise NotImplementedError("TokenCompressedAttention: cross-attention not supported.")

        B, L, D = hidden_states.shape
        H, H_kv, Dh = self.num_heads, self.num_kv_heads, self.head_dim
        device = hidden_states.device
        dtype = hidden_states.dtype

        q = self.q_lin(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
        k = self.k_lin(hidden_states).view(B, L, H_kv, Dh).transpose(1, 2).contiguous()
        v = self.v_lin(hidden_states).view(B, L, H_kv, Dh).transpose(1, 2).contiguous()

        # ---- RoPE ----
        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]
            cos = cos[..., past_len : past_len + L, :]
            sin = sin[..., past_len : past_len + L, :]
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # ---- valid mask ----
        valid_mask_2d = None
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                if attention_mask.dtype == torch.bool:
                    valid_mask_2d = attention_mask
                elif attention_mask.dtype.is_floating_point:
                    valid_mask_2d = attention_mask > -1e4
                else:
                    valid_mask_2d = attention_mask > 0
            elif attention_mask.dim() == 4:
                if attention_mask.dtype == torch.bool:
                    v_mask = attention_mask
                elif attention_mask.dtype.is_floating_point:
                    v_mask = attention_mask > -1e4
                else:
                    v_mask = attention_mask > 0
                valid_mask_2d = v_mask.any(dim=-2).squeeze(1)

        # ---- TC-23: KV-cache mode bypasses compression ----
        in_kv_cache_mode = past_key_value is not None and L == 1
        if in_kv_cache_mode:
            if H_kv != H:
                k = repeat_kv(k, self.num_key_value_groups)
                v = repeat_kv(v, self.num_key_value_groups)
            scale = 1.0 / math.sqrt(Dh)
            scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale
            attn = F.softmax(scores, dim=-1).to(dtype)
            attn = torch.nan_to_num(attn, nan=0.0)
            attn = self.dropout(attn)
            context = torch.matmul(attn, v)
            context = context.transpose(1, 2).contiguous().view(B, L, self.q_head_size)
            if self.is_nested_encoder:
                return context, (attn if output_attentions else None)
            return self.out_lin(context), (attn if output_attentions else None)

        # ---- Token compression ----
        k_comp, v_comp, top_idx, comp_valid = self._select_compressed_tokens(
            hidden_states,
            k,
            v,
            valid_mask_2d,
        )

        if H_kv != H:
            k_comp = repeat_kv(k_comp, self.num_key_value_groups)
            v_comp = repeat_kv(v_comp, self.num_key_value_groups)

        scale = 1.0 / math.sqrt(Dh)
        scores = torch.matmul(q.float(), k_comp.float().transpose(-2, -1).contiguous()) * scale

        # ---- TC-21: causality with fallback ----
        if self.is_decoder:
            q_pos = torch.arange(L, device=device).view(1, 1, L, 1)
            k_orig = top_idx.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, k_keep)
            causal_mask = k_orig > q_pos  # (B, 1, L, k_keep)
            scores = scores.masked_fill(causal_mask, -1e4)

            # TC-21: for queries where all kept tokens are future, force-keep
            # the first kept token (which is CLS = index 0 for decoders,
            # guaranteed by TC-22).
            all_future = causal_mask.all(dim=-1, keepdim=True)  # (B, 1, L, 1)
            if all_future.any():
                # Un-mask the CLS slot for those queries.
                cls_pos = top_idx == 0
                cls_slot = cls_pos.long().argmax(dim=-1, keepdim=True)  # (B, 1)
                cls_slot = cls_slot.unsqueeze(1)  # (B, 1, 1, 1)
                cls_slot = cls_slot.expand(-1, 1, L, 1)
                scores = scores.scatter(-1, cls_slot, 0.0)

        if comp_valid is not None:
            # Never mask out the CLS slot even if comp_valid is False,
            # because TC-22 forced it in.
            cm = comp_valid[:, None, None, :]  # (B, 1, 1, k_keep)
            if self.is_decoder:
                # Un-mask CLS position
                cls_mask = (top_idx == 0)[:, None, None, :]  # (B, 1, 1, k_keep)
                cm = cm | cls_mask
            scores = scores.masked_fill(~cm, -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)
            attn = attn * head_mask.to(attn.dtype)

        context = torch.matmul(attn, v_comp)
        context = context.transpose(1, 2).contiguous().view(B, L, self.q_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_lin(context)

        if output.shape[-1] != D:
            raise RuntimeError(
                f"TokenCompressedAttention: output dim {output.shape[-1]} != hidden_size {D}"
            )

        return output, (attn if output_attentions else None)

Every 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

The type of the None singleton.

var supports_causal

The type of the None singleton.

var supports_gqa

The type of the None singleton.

Instance variables

prop c_proj

Expand source code

@property
def c_proj(self):
    return self.out_lin

prop dense

Expand source code

@property
def dense(self):
    return self.out_lin

prop o_proj

Expand source code

@property
def o_proj(self):
    return self.out_lin

prop out_proj_alias

Expand source code

@property
def out_proj_alias(self):
    return self.out_lin

prop q_lin_alias

Expand source code

@property
def q_lin_alias(self):
    return self.q_lin

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class TokenRouter (hidden_dim: int, keep_ratio: float = 0.75, min_tokens: int = 8)

Expand source code

class TokenRouter(nn.Module):
    r"""Top-k token selector based on learned saliency scores with guaranteed [CLS] retention."""

    def __init__(self, hidden_dim: int, keep_ratio: float = 0.75, min_tokens: int = 8) -> None:
        super().__init__()
        self.net = nn.Sequential(nn.Linear(hidden_dim, 32), nn.GELU(), nn.Linear(32, 1))
        self.keep_ratio = keep_ratio
        self.min_tokens = min_tokens

    def forward(
        self,
        x: torch.Tensor,
        valid_mask: torch.Tensor | None = None,
        is_decoder: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        B, N, D = x.shape
        net_dtype = next(self.net.parameters()).dtype
        scores = torch.sigmoid(self.net(x.to(net_dtype)).squeeze(-1))  # (B, N)
        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask, -10000.0)

        if is_decoder and N > 1:
            causal_mask = torch.tril(torch.ones(N, N, device=x.device, dtype=torch.bool))
            scores = scores.masked_fill(~causal_mask[0], -10000.0)

        # Guarantee preservation of first token (index 0)
        scores[:, 0] = 1e4

        keep = min(
            N,
            max(
                1,
                (
                    min(self.min_tokens, N)
                    if self.min_tokens >= N
                    else max(self.min_tokens, int(N * self.keep_ratio))
                ),
            ),
        )
        idx = scores.topk(keep, dim=1).indices
        idx, _ = torch.sort(idx, dim=1)

        routed = torch.gather(x, 1, idx.unsqueeze(-1).expand(-1, -1, D))
        return routed, idx

Top-k token selector based on learned saliency scores with guaranteed [CLS] retention.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self,
x: torch.Tensor,
valid_mask: torch.Tensor | None = None,
is_decoder: bool = False) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

def forward(
    self,
    x: torch.Tensor,
    valid_mask: torch.Tensor | None = None,
    is_decoder: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    B, N, D = x.shape
    net_dtype = next(self.net.parameters()).dtype
    scores = torch.sigmoid(self.net(x.to(net_dtype)).squeeze(-1))  # (B, N)
    if valid_mask is not None:
        scores = scores.masked_fill(~valid_mask, -10000.0)

    if is_decoder and N > 1:
        causal_mask = torch.tril(torch.ones(N, N, device=x.device, dtype=torch.bool))
        scores = scores.masked_fill(~causal_mask[0], -10000.0)

    # Guarantee preservation of first token (index 0)
    scores[:, 0] = 1e4

    keep = min(
        N,
        max(
            1,
            (
                min(self.min_tokens, N)
                if self.min_tokens >= N
                else max(self.min_tokens, int(N * self.keep_ratio))
            ),
        ),
    )
    idx = scores.topk(keep, dim=1).indices
    idx, _ = torch.sort(idx, dim=1)

    routed = torch.gather(x, 1, idx.unsqueeze(-1).expand(-1, -1, D))
    return routed, idx

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class Trainer (model: transformers.modeling_utils.PreTrainedModel | torch.nn.modules.module.Module | None = None,
args: transformers.training_args.TrainingArguments | None = None,
data_collator: Callable[[list[typing.Any]], dict[str, typing.Any]] | None = None,
train_dataset: Dataset | IterableDataset | datasets.Dataset | None = None,
eval_dataset: Dataset | dict[str, Dataset] | datasets.Dataset | None = None,
processing_class: transformers.tokenization_utils_base.PreTrainedTokenizerBase | transformers.image_processing_utils.BaseImageProcessor | transformers.feature_extraction_utils.FeatureExtractionMixin | transformers.processing_utils.ProcessorMixin | None = None,
model_init: Callable[..., transformers.modeling_utils.PreTrainedModel] | None = None,
compute_loss_func:  | None = None,
compute_metrics: Callable[[transformers.trainer_utils.EvalPrediction], dict] | None = None,
callbacks: list[transformers.trainer_callback.TrainerCallback] | None = None,
optimizers: tuple[torch.optim.optimizer.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None),
optimizer_cls_and_kwargs: tuple[type[torch.optim.optimizer.Optimizer], dict[str, typing.Any]] | None = None,
preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None)

Expand source code

@requires(
    backends=(
        "torch",
        "accelerate",
    )
)
class Trainer:
    """
    Trainer is a simple but feature-complete training and eval loop for PyTorch, optimized for 🤗 Transformers.

    Args:
        model ([`PreTrainedModel`] or `torch.nn.Module`, *optional*):
            The model to train, evaluate or use for predictions. If not provided, a `model_init` must be passed.

            <Tip>

            [`Trainer`] is optimized to work with the [`PreTrainedModel`] provided by the library. You can still use
            your own models defined as `torch.nn.Module` as long as they work the same way as the 🤗 Transformers
            models.

            </Tip>

        args ([`TrainingArguments`], *optional*):
            The arguments to tweak for training. Will default to a basic instance of [`TrainingArguments`] with the
            `output_dir` set to a directory named *tmp_trainer* in the current directory if not provided.
        data_collator (`DataCollator`, *optional*):
            The function to use to form a batch from a list of elements of `train_dataset` or `eval_dataset`. Will
            default to [`default_data_collator`] if no `processing_class` is provided, an instance of
            [`DataCollatorWithPadding`] otherwise if the processing_class is a feature extractor or tokenizer.
        train_dataset (`torch.utils.data.Dataset` | `torch.utils.data.IterableDataset` | `datasets.Dataset`, *optional*):
            The dataset to use for training. If it is a [`~datasets.Dataset`], columns not accepted by the
            `model.forward()` method are automatically removed.

            Note that if it's a `torch.utils.data.IterableDataset` with some randomization and you are training in a
            distributed fashion, your iterable dataset should either use a internal attribute `generator` that is a
            `torch.Generator` for the randomization that must be identical on all processes (and the Trainer will
            manually set the seed of this `generator` at each epoch) or have a `set_epoch()` method that internally
            sets the seed of the RNGs used.
        eval_dataset (`torch.utils.data.Dataset` | dict[str, `torch.utils.data.Dataset`] | `datasets.Dataset`, *optional*):
             The dataset to use for evaluation. If it is a [`~datasets.Dataset`], columns not accepted by the
             `model.forward()` method are automatically removed. If it is a dictionary, it will evaluate on each
             dataset prepending the dictionary key to the metric name.
        processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*):
            Processing class used to process the data. If provided, will be used to automatically process the inputs
            for the model, and it will be saved along the model to make it easier to rerun an interrupted training or
            reuse the fine-tuned model.
        model_init (`Callable[[], PreTrainedModel]`, *optional*):
            A function that instantiates the model to be used. If provided, each call to [`~Trainer.train`] will start
            from a new instance of the model as given by this function.

            The function may have zero argument, or a single one containing the optuna/Ray Tune trial object, to
            be able to choose different architectures according to hyperparameters (such as layer count, sizes of
            inner layers, dropout probabilities etc).
        compute_loss_func (`Callable`, *optional*):
            A function that accepts the raw model outputs, labels, and the number of items in the entire accumulated
            batch (batch_size * gradient_accumulation_steps) and returns the loss. For example, see the default [loss function](https://github.com/huggingface/transformers/blob/052e652d6d53c2b26ffde87e039b723949a53493/src/transformers/trainer.py#L3618) used by [`Trainer`].
        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):
            The function that will be used to compute metrics at evaluation. Must take a [`EvalPrediction`] and return
            a dictionary string to metric values. *Note* When passing TrainingArgs with `batch_eval_metrics` set to
            `True`, your compute_metrics function must take a boolean `compute_result` argument. This will be triggered
            after the last eval batch to signal that the function needs to calculate and return the global summary
            statistics rather than accumulating the batch-level statistics
        callbacks (List of [`TrainerCallback`], *optional*):
            A list of callbacks to customize the training loop. Will add those to the list of default callbacks
            detailed in [here](callback).

            If you want to remove one of the default callbacks used, use the [`Trainer.remove_callback`] method.
        optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`):
            A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your
            model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`.
        optimizer_cls_and_kwargs (`tuple[Type[torch.optim.Optimizer], dict[str, Any]]`, *optional*):
            A tuple containing the optimizer class and keyword arguments to use.
            Overrides `optim` and `optim_args` in `args`. Incompatible with the `optimizers` argument.

            Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before initializing the Trainer.
        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*):
            A function that preprocess the logits right before caching them at each evaluation step. Must take two
            tensors, the logits and the labels, and return the logits once processed as desired. The modifications made
            by this function will be reflected in the predictions received by `compute_metrics`.

            Note that the labels (second parameter) will be `None` if the dataset does not have them.

    Important attributes:

        - **model** -- Always points to the core model. If using a transformers model, it will be a [`PreTrainedModel`]
          subclass.
        - **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the
          original model. This is the model that should be used for the forward pass. For example, under `DeepSpeed`,
          the inner model is wrapped in `DeepSpeed` and then again in `torch.nn.DistributedDataParallel`. If the inner
          model hasn't been wrapped, then `self.model_wrapped` is the same as `self.model`.
        - **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from
          data parallelism, this means some of the model layers are split on different GPUs).
        - **place_model_on_device** -- Whether or not to automatically place the model on the device. Defaults to
          `True` unless model parallel, DeepSpeed, FSDP, full fp16/bf16 eval, or SageMaker MP is active. Can be
          overridden by subclassing `TrainingArguments` and overriding the `place_model_on_device` property.
        - **is_in_train** -- Whether or not a model is currently running `train` (e.g. when `evaluate` is called while
          in `train`)

    """

    # Those methods are not used in Trainer itself but are available as methods for external use.
    from .trainer_pt_utils import (
        get_learning_rates,
        get_num_trainable_parameters,
        get_optimizer_group,
        log_metrics,
        metrics_format,
        save_metrics,
        save_state,
    )

    # ---- Initialization & Validation ----

    def __init__(
        self,
        model: PreTrainedModel | nn.Module | None = None,
        args: TrainingArguments | None = None,
        data_collator: DataCollator | None = None,
        train_dataset: "Dataset | IterableDataset | datasets.Dataset | None" = None,
        eval_dataset: "Dataset | dict[str, Dataset] | datasets.Dataset | None" = None,
        processing_class: PreTrainedTokenizerBase
        | BaseImageProcessor
        | FeatureExtractionMixin
        | ProcessorMixin
        | None = None,
        model_init: Callable[..., PreTrainedModel] | None = None,
        compute_loss_func: Callable | None = None,
        compute_metrics: Callable[[EvalPrediction], dict] | None = None,
        callbacks: list[TrainerCallback] | None = None,
        optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None),
        optimizer_cls_and_kwargs: tuple[type[torch.optim.Optimizer], dict[str, Any]] | None = None,
        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
    ):
        # Init flow:
        #   1. Args & seed               – defaults, determinism
        #   2. Accelerator & logging     – accelerator, memory tracker, log level, device setup
        #   3. Model resolution          – model / model_init, Liger Kernel, quantization checks
        #   4. Distributed strategy      – model-parallel, FSDP, SageMaker MP flags
        #   5. Device placement          – move model to device, model wrapping
        #   6. Model introspection       – loss kwargs, label names, label smoother
        #   7. Store init arguments      – data, callables, optimizer, scheduler, validation
        #   8. Callbacks                 – reporting integrations, JIT checkpoint, progress bar
        #   9. Hub & output              – repo init, output directory
        #  10. Training state            – TrainerState, TrainerControl, internal bookkeeping
        #  11. Finalize                  – use_cache, XLA FSDPv2 mesh, memory tracker stop

        # ---- 1. Args & seed --------------------------------------------------------
        if args is None:
            output_dir = "tmp_trainer"
            logger.info(f"No `TrainingArguments` passed, using `output_dir={output_dir}`.")
            args = TrainingArguments(output_dir=output_dir)
        self.args = args
        # Seed must be set before instantiating the model when using model_init
        enable_full_determinism(self.args.seed) if self.args.full_determinism else set_seed(self.args.seed)

        # ---- 2. Accelerator & logging ----------------------------------------------
        # `create_accelerator_and_postprocess` reads self.model and self.args,
        # and may set self.deepspeed — store temporary refs before calling it.
        self.deepspeed = None
        self.model = model
        self.create_accelerator_and_postprocess()

        self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics)
        self._memory_tracker.start()

        log_level = args.get_process_log_level()
        logging.set_verbosity(log_level)

        args._setup_devices  # force device and distributed setup init explicitly

        # ---- 3. Model resolution ----------------------------------------------------
        if model is None:
            if model_init is not None:
                self.model_init = model_init
                model = self.call_model_init()
            else:
                raise RuntimeError("`Trainer` requires either a `model` or `model_init` argument")
        else:
            if model_init is not None:
                raise ValueError("`Trainer` requires either a `model` or `model_init` argument, but not both.")
            self.model_init = model_init

        if model.__class__.__name__ in MODEL_MAPPING_NAMES:
            raise ValueError(
                f"The model you have picked ({model.__class__.__name__}) cannot be used as is for training: it only "
                "computes hidden states and does not accept any labels. You should choose a model with a head "
                "suitable for your task like any of the `AutoModelForXxx` listed at "
                "https://huggingface.co/docs/transformers/model_doc/auto"
            )

        validate_quantization_for_training(model)

        # ---- 4. Distributed strategy ------------------------------------------------
        self.is_model_parallel = False
        if getattr(model, "hf_device_map", None) is not None:
            devices = [device for device in set(model.hf_device_map.values()) if device not in ["cpu", "disk"]]
            if len(devices) > 1:
                self.is_model_parallel = True
            elif len(devices) == 1:
                self.is_model_parallel = self.args.device != torch.device(devices[0])

        self.is_fsdp_xla_enabled = args.fsdp_config["xla"]
        if len(args.fsdp) > 0:
            if self.is_deepspeed_enabled:
                raise ValueError(
                    "Using --fsdp xxx together with --deepspeed is not possible, deactivate one of those flags."
                )
            if not args.fsdp_config["xla"] and args.parallel_mode != ParallelMode.DISTRIBUTED:
                raise ValueError("Using fsdp only works in distributed training.")

        # Postpone switching model to cuda when MP, DeepSpeed, full bf16/fp16 eval, or FSDP
        if args.place_model_on_device is not None:
            self.place_model_on_device = args.place_model_on_device
        elif (
            self.is_model_parallel
            or self.is_deepspeed_enabled
            or (args.fp16_full_eval or args.bf16_full_eval)
            or self.is_fsdp_xla_enabled
            or self.is_fsdp_enabled
            or is_sagemaker_mp_enabled()
        ):
            self.place_model_on_device = False
        else:
            self.place_model_on_device = True

        # ---- 5. Device placement ----------------------------------------------------
        # Bnb Quantized models don't support `.to` operation.
        if (
            self.place_model_on_device
            and getattr(model, "quantization_method", None) != QuantizationMethod.BITS_AND_BYTES
        ):
            self._move_model_to_device(model, args.device)

        # Force n_gpu to 1 to avoid DataParallel as MP will manage the GPUs
        if self.is_model_parallel:
            self.args._n_gpu = 1

        # `self.model is self.model_wrapped` is used later to check if it's wrapped
        self.model_wrapped = model
        self.model = model

        # ---- 6. Model introspection -------------------------------------------------
        unwrapped_model = unwrap_peft_model(self.accelerator.unwrap_model(model))

        if hasattr(unwrapped_model, "accepts_loss_kwargs"):
            self.model_accepts_loss_kwargs = unwrapped_model.accepts_loss_kwargs
        else:
            forward_params = inspect.signature(unwrapped_model.forward).parameters
            self.model_accepts_loss_kwargs = any(
                k.kind == inspect.Parameter.VAR_KEYWORD for k in forward_params.values()
            )

        # Sequence Parallelism computes its own good_tokens count
        pc = getattr(self.accelerator, "parallelism_config", None)
        if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled:
            self.model_accepts_loss_kwargs = False

        model_to_inspect = unwrap_peft_model(self.model)
        default_label_names = find_labels(model_to_inspect.__class__)
        self.label_names = default_label_names if self.args.label_names is None else self.args.label_names
        self.can_return_loss = can_return_loss(model_to_inspect.__class__)

        if self.args.label_smoothing_factor != 0:
            if getattr(self.model.config, "problem_type", None) == "multi_label_classification":
                warnings.warn(
                    "Label smoothing is not compatible with multi-label classification. "
                    "Disabling label smoothing for this training run.",
                    UserWarning,
                )
                self.label_smoother = None
            else:
                self.label_smoother = LabelSmoother(epsilon=self.args.label_smoothing_factor)
        else:
            self.label_smoother = None

        # ---- 7. Store init arguments ------------------------------------------------
        # Data
        default_collator = (
            DataCollatorWithPadding(processing_class)
            if processing_class is not None
            and isinstance(processing_class, (PreTrainedTokenizerBase, SequenceFeatureExtractor))
            else default_data_collator
        )
        self.data_collator = data_collator if data_collator is not None else default_collator
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.processing_class = processing_class
        self.neftune_noise_alpha = args.neftune_noise_alpha

        # Callables
        self.compute_loss_func = compute_loss_func
        self.compute_metrics = compute_metrics
        self.preprocess_logits_for_metrics = preprocess_logits_for_metrics

        # Optimizer & scheduler
        self.optimizer, self.lr_scheduler = optimizers
        self.optimizer_cls_and_kwargs = optimizer_cls_and_kwargs

        self._validate_args()

        # ---- 8. Callbacks -----------------------------------------------------------
        default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)

        if self.args.enable_jit_checkpoint:
            from .trainer_jit_checkpoint import JITCheckpointCallback

            jit_callback = JITCheckpointCallback()
            default_callbacks = default_callbacks + [jit_callback]
            jit_callback.set_trainer(self)

        callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
        self.callback_handler = CallbackHandler(
            callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler
        )
        self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK)

        # ---- 9. Hub & output ---------------------------------------------------------
        self.hub_model_id = None  # Set by init_hf_repo() when push_to_hub is enabled
        if self.args.push_to_hub:
            self.init_hf_repo()
        if self.args.should_save:
            os.makedirs(self.args.output_dir, exist_ok=True)

        # ---- 10. Training state -----------------------------------------------------
        self.control = TrainerControl()
        self.state = TrainerState(
            is_local_process_zero=self.is_local_process_zero(),
            is_world_process_zero=self.is_world_process_zero(),
            stateful_callbacks=[
                cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
            ],
        )
        self.is_in_train = False  # True between train() entry and exit
        self.hp_name = None  # Set by hyperparameter_search() to label the trial
        self.hp_search_backend = None  # Set by hyperparameter_search() (optuna / ray / wandb)
        # Per-process FLOP counter; accumulated into self.state.total_flos then reset
        self.current_flos = 0
        # Set True by _setup_loggers() on first call to self.log()
        self._loggers_initialized = False
        # Lazily filled by _set_signature_columns_if_needed(); caches model.forward param names
        self._signature_columns = None
        # Effective batch size; may be reduced by find_executable_batch_size
        self._train_batch_size = args.train_batch_size
        # Guards one-time LR scheduler creation in create_optimizer_and_scheduler
        self._created_lr_scheduler = False

        self.control = self.callback_handler.on_init_end(self.args, self.state, self.control)

        # ---- 11. Finalize -----------------------------------------------------------
        if getattr(self.model, "config", None) is not None:
            self.model.config.use_cache = self.args.use_cache

        self.is_fsdp_xla_v2_enabled = args.fsdp_config.get("xla_fsdp_v2", False)
        if self.is_fsdp_xla_v2_enabled:
            if not IS_XLA_FSDPV2_POST_2_2:
                raise ValueError("FSDPv2 requires `torch_xla` 2.2 or higher.")
            num_devices = xr.global_runtime_device_count()
            xs.set_global_mesh(xs.Mesh(np.array(range(num_devices)), (num_devices, 1), axis_names=("fsdp", "tensor")))
        self.is_fsdp_xla_v1_enabled = self.is_fsdp_xla_enabled and not self.is_fsdp_xla_v2_enabled

        self._memory_tracker.stop_and_update_metrics()

    def _validate_args(self) -> None:
        """Validate constructor arguments and fail fast on incompatible combinations."""
        args = self.args

        # --- SageMaker Model Parallel mixed-precision validation ---
        if is_sagemaker_mp_enabled():
            if args.bf16:
                raise ValueError("SageMaker Model Parallelism does not support BF16 yet. Please use FP16 instead ")
            if args.fp16 != smp.state.cfg.fp16:
                logger.warning(
                    f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, "
                    f"but FP16 provided in trainer argument is {args.fp16}, "
                    f"setting to {smp.state.cfg.fp16}"
                )
                args.fp16 = smp.state.cfg.fp16

        # --- Training-argument validations ---
        if args.batch_eval_metrics and self.compute_metrics is not None:
            if "compute_result" not in inspect.signature(self.compute_metrics).parameters:
                raise ValueError(
                    "When using `batch_eval_metrics`, your `compute_metrics` function must take a `compute_result`"
                    " boolean argument which will be triggered after the last batch of the eval set to signal that the"
                    " summary statistics should be returned by the function."
                )
        if args.eval_strategy is not None and args.eval_strategy != "no" and self.eval_dataset is None:
            raise ValueError(
                f"You have set `args.eval_strategy` to {args.eval_strategy} but you didn't pass an `eval_dataset` to `Trainer`. Either set `args.eval_strategy` to `no` or pass an `eval_dataset`. "
            )
        if args.save_strategy == SaveStrategy.BEST or args.load_best_model_at_end:
            if args.metric_for_best_model is None:
                raise ValueError(
                    "`args.metric_for_best_model` must be provided when using 'best' save_strategy or if `args.load_best_model_at_end` is set to `True`."
                )

        # --- Optimizer validations ---
        if self.optimizer_cls_and_kwargs is not None and self.optimizer is not None:
            raise RuntimeError("Passing both `optimizers` and `optimizer_cls_and_kwargs` arguments is incompatible.")
        if self.model_init is not None and (self.optimizer is not None or self.lr_scheduler is not None):
            raise RuntimeError(
                "Passing a `model_init` is incompatible with providing the `optimizers` argument. "
                "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
            )
        if is_torch_xla_available() and self.optimizer is not None:
            for param in self.model.parameters():
                model_device = param.device
                break
            for param_group in self.optimizer.param_groups:
                if len(param_group["params"]) > 0:
                    optimizer_device = param_group["params"][0].device
                    break
            if model_device != optimizer_device:
                raise ValueError(
                    "The model and the optimizer parameters are not on the same device, which probably means you"
                    " created an optimizer around your model **before** putting on the device and passing it to the"
                    " `Trainer`. Make sure the lines `import torch_xla.core.xla_model as xm` and"
                    " `model.to(xm.xla_device())` is performed before the optimizer creation in your script."
                )
        if (self.is_fsdp_xla_enabled or self.is_fsdp_enabled) and (
            self.optimizer is not None or self.lr_scheduler is not None
        ):
            raise RuntimeError(
                "Passing `optimizers` is not allowed if PyTorch FSDP is enabled. "
                "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
            )

        # --- Dataset validations ---
        if not callable(self.data_collator) and callable(getattr(self.data_collator, "collate_batch", None)):
            raise TypeError("The `data_collator` should be a simple callable (function, class with `__call__`).")
        if args.max_steps > 0 and args.num_train_epochs > 0:
            logger.info("max_steps is given, it will override any value given in num_train_epochs")
        if self.train_dataset is not None and not has_length(self.train_dataset) and args.max_steps <= 0:
            raise ValueError(
                "The train_dataset does not implement __len__, max_steps has to be specified. "
                "The number of steps needs to be known in advance for the learning rate scheduler."
            )

        if self.train_dataset is not None and isinstance(self.train_dataset, torch.utils.data.IterableDataset):
            logger.info(
                f"The `train_sampling_strategy='{args.train_sampling_strategy}'` option is ignored when using an `IterableDataset`. "
                "Samplers cannot be used with IterableDataset as they require indexed access to the dataset."
            )

    def _build_accelerator_args(self, **kwargs) -> dict[str, Any]:
        """Helper method to build accelerator-specific keyword arguments."""
        args = {
            "mixed_precision": self.args.mixed_precision,
            "deepspeed_plugin": self.args.deepspeed_plugin,
        }
        args.update(kwargs)

        if self.args.ddp_find_unused_parameters is not None:
            find_unused = self.args.ddp_find_unused_parameters
        elif isinstance(self.model, PreTrainedModel):
            # find_unused_parameters breaks checkpointing as per
            # https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021
            find_unused = not (self.model.is_gradient_checkpointing or self.args.gradient_checkpointing)
        else:
            find_unused = True

        ddp_kwargs = {"find_unused_parameters": find_unused}
        if self.args.ddp_bucket_cap_mb is not None:
            ddp_kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb
        if self.args.ddp_broadcast_buffers is not None:
            ddp_kwargs["broadcast_buffers"] = self.args.ddp_broadcast_buffers
        if self.args.ddp_static_graph is not None:
            ddp_kwargs["static_graph"] = self.args.ddp_static_graph

        args["kwargs_handlers"] = [DistributedDataParallelKwargs(**ddp_kwargs)]

        # We defer compatibility checks to accelerator
        if self.args.parallelism_config is not None:
            min_accelerate_version = "1.12.0"
            if not is_accelerate_available(min_accelerate_version):
                raise ImportError(
                    f"ParallelismConfig requires accelerate>={min_accelerate_version}). Please upgrade accelerate to use this feature."
                )
            args["parallelism_config"] = self.args.parallelism_config

        if getattr(self.model, "tp_size", None) is not None and self.model.tp_size > 1:
            if self.args.parallelism_config is None:
                if is_accelerate_available("1.12.0"):
                    if self.args.parallelism_config is None:
                        from accelerate import ParallelismConfig

                        args["parallelism_config"] = ParallelismConfig(tp_size=self.model.tp_size)
                else:
                    raise ValueError("Requires accelerate>1.12.0 to use Tensor Parallelism.")
            elif args["parallelism_config"].tp_size != self.model.tp_size:
                args["parallelism_config"].tp_size = self.model.tp_size

        if is_accelerate_available("1.2.0"):
            # it we don't have the correct version, we will rely on env var instead that were set in TrainingArguments
            from accelerate.utils import TorchDynamoPlugin

            dynamo_plugin = TorchDynamoPlugin(
                backend=self.args.torch_compile_backend, mode=self.args.torch_compile_mode
            )
            args["dynamo_plugin"] = dynamo_plugin

        return args

    def create_accelerator_and_postprocess(self) -> None:
        """Create the accelerator and perform post-creation setup (FSDP, DeepSpeed, etc.)."""
        # We explicitly don't rely on the `Accelerator` to do gradient accumulation
        grad_acc_kwargs = {}
        if self.args.accelerator_config.gradient_accumulation_kwargs is not None:
            grad_acc_kwargs = self.args.accelerator_config.gradient_accumulation_kwargs

        # check if num_steps is attempted to be passed in gradient_accumulation_kwargs
        if "num_steps" in grad_acc_kwargs:
            if self.args.gradient_accumulation_steps > 1:
                # raise because we do not know which setting is intended.
                raise ValueError(
                    "The `AcceleratorConfig`'s `num_steps` is set but `gradient_accumulation_steps` is greater than 1 in the passed `TrainingArguments`"
                    "If using the passed `AcceleratorConfig` is desired, do not set the `TrainingArguments` `gradient_accumulation_steps`."
                )
            else:
                self.args.gradient_accumulation_steps = grad_acc_kwargs["num_steps"]

        # The Trainer handles GAS itself, so GAS=1 in Accelerate to avoid any double-division
        grad_acc_kwargs["num_steps"] = 1

        # Just making sure that gradient_state have the correct values passed.
        # We don't rely on `accumulate` from accelerate to set sync_gradients in gradient_state.
        # Rather, we do it ourselves by setting self.accelerator.gradient_state._set_sync_gradients.
        gradient_accumulation_plugin = GradientAccumulationPlugin(**grad_acc_kwargs)

        accelerator_config = self.args.accelerator_config.to_dict()

        # Extract dataloader config params from accelerator config
        dataloader_params = ["split_batches", "dispatch_batches", "even_batches", "use_seedable_sampler"]
        dataloader_config = DataLoaderConfiguration(
            **{param: accelerator_config.pop(param) for param in dataloader_params}
        )
        dataloader_config.data_seed = self.args.data_seed

        non_blocking = accelerator_config.pop("non_blocking")

        if non_blocking and not self.args.dataloader_pin_memory:
            logger.warning(
                "`non_blocking` is enabled but `dataloader_pin_memory` is not. For the best performance, it's recommended to enable both."
            )
        dataloader_config.non_blocking = non_blocking
        # this would have been updated above, no need for it anymore
        accelerator_config.pop("gradient_accumulation_kwargs")

        fsdp_plugin = None
        if self.args.fsdp_plugin_args is not None:
            from accelerate.utils import FullyShardedDataParallelPlugin

            fsdp_plugin = FullyShardedDataParallelPlugin(**self.args.fsdp_plugin_args)

        args = self._build_accelerator_args(
            dataloader_config=dataloader_config,
            fsdp_plugin=fsdp_plugin,
            gradient_accumulation_plugin=gradient_accumulation_plugin,
        )

        # create accelerator object
        self.accelerator = Accelerator(**args)
        # some Trainer classes need to use `gather` instead of `gather_for_metrics`, thus we store a flag
        self.gather_function = self.accelerator.gather_for_metrics

        if "use_gather_object" in inspect.signature(self.gather_function).parameters:
            self.gather_function = functools.partial(
                self.gather_function, use_gather_object=self.args.eval_use_gather_object
            )

        # deepspeed and accelerate flags covering both trainer args and accelerate launcher
        self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None
        self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None

        # post accelerator creation setup
        if self.is_fsdp_enabled:
            fsdp_plugin = self.accelerator.state.fsdp_plugin
            for param in ["limit_all_gathers", "activation_checkpointing"]:
                setattr(fsdp_plugin, param, self.args.fsdp_config.get(param, getattr(fsdp_plugin, param)))
            if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing:
                raise ValueError(
                    "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg "
                    "can't be set to True simultaneously. Please use FSDP's activation_checkpointing logic "
                    "when using FSDP."
                )

        if self.is_deepspeed_enabled and getattr(self.args, "hf_deepspeed_config", None) is None:
            propagate_args_to_deepspeed(self.accelerator, self.args)

        # `save_only_model` can't be used with DeepSpeed/FSDP along with `load_best_model_at_end`
        if (
            self.args.save_only_model
            and (self.is_deepspeed_enabled or self.is_fsdp_enabled)
            and self.args.load_best_model_at_end
        ):
            wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP"
            raise ValueError(f"{wrapper} can't be used with `save_only_model` along with `load_best_model_at_end`.")

        # `auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3
        if (
            self.is_deepspeed_enabled
            and self.accelerator.state.deepspeed_plugin.zero_stage == 3
            and self.args.auto_find_batch_size
        ):
            raise ValueError(
                "`auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3. Please consider using Zero-2, Zero-1, or FSDP"
            )
        if (
            self.args.save_only_model
            and self.is_fsdp_enabled
            and "SHARDED_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type)
        ):
            raise ValueError("save_only_model option is not compatible with FSDP state dict type 'SHARDED_STATE_DICT'")

    # ---- Data Loading ----

    def get_train_dataloader(self) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")

        return self._get_dataloader(
            dataset=self.train_dataset,
            description="Training",
            batch_size=self._train_batch_size,
            sampler_fn=self._get_train_sampler,
            is_training=True,
        )

    def get_eval_dataloader(self, eval_dataset: str | Dataset | None = None) -> DataLoader:
        """
        Returns the evaluation [`~torch.utils.data.DataLoader`].

        Subclass and override this method if you want to inject some custom behavior.

        Args:
            eval_dataset (`str` or `torch.utils.data.Dataset`, *optional*):
                If a `str`, will use `self.eval_dataset[eval_dataset]` as the evaluation dataset. If a `Dataset`, will override `self.eval_dataset` and must implement `__len__`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed.
        """
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")

        # If we have persistent workers, don't do a fork bomb especially as eval datasets
        # don't change during training
        dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval"
        if (
            hasattr(self, "_eval_dataloaders")
            and dataloader_key in self._eval_dataloaders
            and self.args.dataloader_persistent_workers
        ):
            return self._eval_dataloaders[dataloader_key]

        eval_dataset = (
            self.eval_dataset[eval_dataset]
            if isinstance(eval_dataset, str)
            else eval_dataset
            if eval_dataset is not None
            else self.eval_dataset
        )

        return self._get_dataloader(
            dataset=eval_dataset,
            description="Evaluation",
            batch_size=self.args.eval_batch_size,
            sampler_fn=self._get_eval_sampler,
            dataloader_key=dataloader_key,
        )

    def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
        """
        Returns the test [`~torch.utils.data.DataLoader`].

        Subclass and override this method if you want to inject some custom behavior.

        Args:
            test_dataset (`torch.utils.data.Dataset`, *optional*):
                The test dataset to use. If it is a [`~datasets.Dataset`], columns not accepted by the
                `model.forward()` method are automatically removed. It must implement `__len__`.
        """
        return self._get_dataloader(
            dataset=test_dataset,
            description="test",
            batch_size=self.args.eval_batch_size,
            sampler_fn=self._get_eval_sampler,
        )

    def num_examples(self, dataloader: DataLoader) -> int:
        """
        Helper to get number of samples in a [`~torch.utils.data.DataLoader`] by accessing its dataset. When
        dataloader.dataset does not exist or has no length, estimates as best it can
        """
        try:
            dataset = dataloader.dataset
            # Special case for IterableDatasetShard, we need to dig deeper
            if isinstance(dataset, IterableDatasetShard):
                return len(dataloader.dataset.dataset)
            return len(dataloader.dataset)
        except (NameError, AttributeError, TypeError):  # no dataset or length, estimate by length of dataloader
            return len(dataloader) * self.args.per_device_train_batch_size

    def _get_dataloader(
        self,
        dataset: Dataset,
        description: str,
        batch_size: int,
        sampler_fn: Callable[[Dataset], torch.utils.data.Sampler] | None = None,
        is_training: bool = False,
        dataloader_key: str | None = None,
    ) -> DataLoader:
        """Create a [`~torch.utils.data.DataLoader`] from the given dataset."""

        data_collator = self.data_collator
        if is_datasets_available() and isinstance(dataset, datasets.Dataset):
            dataset = self._remove_unused_columns(dataset, description=description)
        else:
            data_collator = self._get_collator_with_removed_columns(self.data_collator, description=description)

        # MPS requrires forking if multiple workers are specified
        should_fork = torch.backends.mps.is_available() and self.args.dataloader_num_workers > 1

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
            "multiprocessing_context": "fork" if should_fork else None,
        }

        if not isinstance(dataset, torch.utils.data.IterableDataset):
            if sampler_fn is not None:
                dataloader_params["sampler"] = sampler_fn(dataset)
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
            if is_training:
                dataloader_params["worker_init_fn"] = partial(
                    seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
                )

        dataloader = self.accelerator.prepare(DataLoader(dataset, **dataloader_params))

        # Store the prepared dataloader for subsequent evaluations if using persistent workers.
        if dataloader_key is not None and self.args.dataloader_persistent_workers:
            if hasattr(self, "_eval_dataloaders"):
                self._eval_dataloaders[dataloader_key] = dataloader
            else:
                self._eval_dataloaders = {dataloader_key: dataloader}

        return dataloader

    def _get_train_sampler(self, train_dataset: Dataset | None = None) -> torch.utils.data.Sampler | None:
        """Return the training sampler based on `train_sampling_strategy`."""
        if train_dataset is None:
            train_dataset = self.train_dataset
        if train_dataset is None or not has_length(train_dataset):
            return None

        # Build the sampler.
        if self.args.train_sampling_strategy == "group_by_length":
            if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
                lengths = (
                    train_dataset[self.args.length_column_name]
                    if self.args.length_column_name in train_dataset.column_names
                    else None
                )
            else:
                lengths = None
            model_input_name = (
                self.processing_class.model_input_names[0] if self.processing_class is not None else None
            )
            return LengthGroupedSampler(
                self.args.train_batch_size * self.args.gradient_accumulation_steps,
                dataset=train_dataset,
                lengths=lengths,
                model_input_name=model_input_name,
            )
        elif self.args.train_sampling_strategy == "sequential":
            return SequentialSampler(train_dataset)
        else:
            return RandomSampler(train_dataset)

    def _get_eval_sampler(self, eval_dataset: Dataset) -> torch.utils.data.Sampler | None:
        """Return the evaluation sampler, using sequential ordering when not distributed."""
        if eval_dataset is None or not has_length(eval_dataset):
            return None

        if self.args.train_sampling_strategy == "group_by_length":
            if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
                lengths = (
                    eval_dataset[self.args.length_column_name]
                    if self.args.length_column_name in eval_dataset.column_names
                    else None
                )
            else:
                lengths = None
            model_input_name = (
                self.processing_class.model_input_names[0] if self.processing_class is not None else None
            )
            return LengthGroupedSampler(
                self.args.eval_batch_size,
                dataset=eval_dataset,
                lengths=lengths,
                model_input_name=model_input_name,
            )

        if self.args.world_size <= 1:
            return SequentialSampler(eval_dataset)
        else:
            return None

    def _set_signature_columns_if_needed(self) -> None:
        """Populate `_signature_columns` from the model's forward signature if not already set."""
        if self._signature_columns is None:
            # Inspect model forward signature to keep only the arguments it accepts.
            model_to_inspect = self.model
            if _is_peft_model(self.model):
                if hasattr(self.model, "get_base_model"):
                    model_to_inspect = self.model.get_base_model()
                else:
                    # PeftMixedModel do not provide a `get_base_model` method
                    model_to_inspect = self.model.base_model.model
            signature = inspect.signature(model_to_inspect.forward)
            self._signature_columns = list(signature.parameters.keys())
            # Labels may be named label or label_ids, the default data collator handles that.
            self._signature_columns += list(set(["label", "label_ids"] + self.label_names))

    def _remove_unused_columns(
        self, dataset: "datasets.Dataset", description: str | None = None
    ) -> "datasets.Dataset":
        """Remove dataset columns not accepted by the model's forward method."""
        if not self.args.remove_unused_columns:
            return dataset
        self._set_signature_columns_if_needed()
        signature_columns = self._signature_columns

        ignored_columns = list(set(dataset.column_names) - set(signature_columns))
        if len(ignored_columns) > 0:
            dset_description = "" if description is None else f"in the {description} set"
            logger.info(
                f"The following columns {dset_description} don't have a corresponding argument in "
                f"`{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}."
                f" If {', '.join(ignored_columns)} are not expected by `{self.model.__class__.__name__}.forward`, "
                " you can safely ignore this message."
            )

        columns = [k for k in signature_columns if k in dataset.column_names]
        if len(columns) == 0:
            raise ValueError(
                f"No columns in the dataset match the model's forward method signature: ({', '.join(signature_columns)}). "
                f"The following columns have been ignored: [{', '.join(ignored_columns)}]. "
                "Please check the dataset and model. You may need to set `remove_unused_columns=False` in `TrainingArguments`."
            )

        if version.parse(datasets.__version__) < version.parse("1.4.0"):
            dataset.set_format(
                type=dataset.format["type"], columns=columns, format_kwargs=dataset.format["format_kwargs"]
            )
            return dataset
        else:
            return dataset.remove_columns(ignored_columns)

    def _get_collator_with_removed_columns(self, data_collator: Callable, description: str | None = None) -> Callable:
        """Wrap the data collator in a callable removing unused columns."""
        if not self.args.remove_unused_columns:
            return data_collator
        self._set_signature_columns_if_needed()
        signature_columns = self._signature_columns

        remove_columns_collator = RemoveColumnsCollator(
            data_collator=data_collator,
            signature_columns=signature_columns,
            logger=logger,
            description=description,
            model_name=self.model.__class__.__name__,
        )
        return remove_columns_collator

    # ---- Optimizer & Scheduler & Learning rate ----

    def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
        """
        Setup the optimizer and the learning rate scheduler.

        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method (or `create_optimizer` and/or
        `create_scheduler`) in a subclass.
        """
        self.create_optimizer()
        self.create_scheduler(num_training_steps=num_training_steps)

    def create_optimizer(self, model=None) -> torch.optim.Optimizer:
        """
        Setup the optimizer.

        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method in a subclass.

        Returns:
            `torch.optim.Optimizer`: The optimizer instance.
        """
        opt_model = self.model if model is None else model

        if self.optimizer is None:
            decay_parameters = self.get_decay_parameter_names(opt_model)
            optimizer_grouped_parameters = [
                {
                    "params": [
                        p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                    ],
                    "weight_decay": self.args.weight_decay,
                },
                {
                    "params": [
                        p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                    ],
                    "weight_decay": 0.0,
                },
            ]

            if self.optimizer_cls_and_kwargs is not None:
                optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
            else:
                optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, opt_model)

            # Check if this is a factory (for complex optimizers like Muon, Dion)
            # Factories are instantiated first, then called with (opt_model, **kwargs)
            if is_optimizer_factory(optimizer_cls):
                self.optimizer = optimizer_cls()(opt_model, **optimizer_kwargs)
            else:
                # Standard optimizer class instantiation
                # Overwrite `params` in case it's created by `get_optimizer_cls_and_kwargs`
                # e.g. for GaLore optimizer.
                if "params" in optimizer_kwargs:
                    optimizer_grouped_parameters = optimizer_kwargs.pop("params")

                # Overwrite `model` in case it's created by `get_optimizer_cls_and_kwargs`
                # e.g. for LOMO optimizer.
                if "model" in optimizer_kwargs:
                    optimizer_grouped_parameters = optimizer_kwargs.pop("model")

                # For layer-wise dummy optimizers we overwrite optimizer_grouped_parameters with `optimizer_dict`
                # to avoid arguments conflicts.
                if "optimizer_dict" in optimizer_kwargs:
                    optimizer_grouped_parameters = optimizer_kwargs.pop("optimizer_dict")

                self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

            if "bitsandbytes" in str(optimizer_cls) and optimizer_kwargs.get("optim_bits", None) == 8:
                import bitsandbytes

                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()

                skipped = 0
                for module in opt_model.modules():
                    if isinstance(module, nn.Embedding):
                        skipped += sum({p.data_ptr(): p.numel() for p in module.parameters()}.values())
                        logger.info(f"skipped {module}: {skipped / 2**20}M params")
                        manager.register_module_override(module, "weight", {"optim_bits": 32})
                        logger.debug(f"bitsandbytes: will optimize {module} in fp32")
                logger.info(f"skipped: {skipped / 2**20}M params")

        if is_sagemaker_mp_enabled():
            self.optimizer = smp.DistributedOptimizer(self.optimizer)

        return self.optimizer

    def create_scheduler(
        self, num_training_steps: int, optimizer: torch.optim.Optimizer | None = None
    ) -> torch.optim.lr_scheduler.LRScheduler:
        """
        Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or
        passed as an argument.

        Args:
            num_training_steps (int): The number of training steps to do.

        Returns:
            `torch.optim.lr_scheduler.LRScheduler`: The learning rate scheduler instance.
        """
        if self.lr_scheduler is None:
            if optimizer is None:
                if is_sagemaker_mp_enabled() and smp.state.cfg.fp16:
                    # If fp16 is enabled, we unwrap the optimizer
                    optimizer = self.optimizer.optimizer
                else:
                    optimizer = self.optimizer
            self.lr_scheduler = get_scheduler(
                self.args.lr_scheduler_type,
                optimizer=optimizer,
                num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
                num_training_steps=num_training_steps,
                scheduler_specific_kwargs=self.args.lr_scheduler_kwargs,
            )
            self._created_lr_scheduler = True
        return self.lr_scheduler

    @staticmethod
    def get_optimizer_cls_and_kwargs(args: TrainingArguments, model: PreTrainedModel | None = None) -> tuple[Any, Any]:
        """
        Returns the optimizer class and optimizer parameters based on the training arguments.

        Args:
            args (`transformers.training_args.TrainingArguments`):
                The training arguments for the training session.
            model (`PreTrainedModel`, *optional*):
                The model being trained. Required for some optimizers (GaLore, Apollo, LOMO).

        Returns:
            A tuple containing the optimizer class and a dictionary of optimizer keyword arguments.
        """
        ctx = OptimizerContext(
            args=args,
            model=model,
            optimizer_kwargs={"lr": args.learning_rate},
            adam_kwargs={
                "betas": (args.adam_beta1, args.adam_beta2),
                "eps": args.adam_epsilon,
            },
            optim_args=_parse_optim_args(args.optim_args),
        )

        handler = _OPTIMIZER_HANDLERS.get(args.optim)
        if handler is None:
            raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {args.optim}")

        return handler(ctx)

    def get_decay_parameter_names(self, model: nn.Module) -> list[str]:
        """
        Get all parameter names that weight decay will be applied to.

        This function filters out parameters in two ways:
        1. By layer type (instances of layers specified in ALL_LAYERNORM_LAYERS)
        2. By parameter name patterns (containing 'bias', or variation of 'norm')
        """
        forbidden_name_patterns = [r"bias", r"layernorm", r"rmsnorm", r"(?:^|\.)norm(?:$|\.)", r"_norm(?:$|\.)"]
        decay_parameters = get_parameter_names(model, [nn.LayerNorm], forbidden_name_patterns)
        return decay_parameters

    def _get_learning_rate(self) -> float:
        """
        Returns the current learning rate from the scheduler.

        Handles DeepSpeed's dynamic loss scaling warmup period where `get_last_lr` may fail.
        """
        if self.is_deepspeed_enabled:
            # with deepspeed's fp16 and dynamic loss scale enabled the optimizer/scheduler steps may
            # not run for the first few dozen steps while loss scale is too large, and thus during
            # that time `get_last_lr` will fail if called during that warm up stage, so work around it:
            try:
                last_lr = self.lr_scheduler.get_last_lr()[0]
            except AssertionError as e:
                if "need to call step" in str(e):
                    logger.warning("tried to get lr value before scheduler/optimizer started stepping, returning lr=0")
                    last_lr = 0
                else:
                    raise
        else:
            if isinstance(self.lr_scheduler, (torch.optim.lr_scheduler.ReduceLROnPlateau, GreedyLR)):
                last_lr = self.optimizer.param_groups[0]["lr"]
            else:
                last_lr = self.lr_scheduler.get_last_lr()[0]

        if torch.is_tensor(last_lr):
            last_lr = last_lr.item()
        return last_lr

    # ---- Training ----

    def train(
        self,
        resume_from_checkpoint: str | bool | None = None,
        trial: "optuna.Trial | dict[str, Any] | None" = None,
        ignore_keys_for_eval: list[str] | None = None,
    ) -> TrainOutput:
        """
        Main training entry point.

        Args:
            resume_from_checkpoint (`str` or `bool`, *optional*):
                If a `str`, local path to a saved checkpoint as saved by a previous instance of [`Trainer`]. If a
                `bool` and equals `True`, load the last checkpoint in *args.output_dir* as saved by a previous instance
                of [`Trainer`]. If present, training will resume from the model/optimizer/scheduler states loaded here.
            trial (`optuna.Trial` or `dict[str, Any]`, *optional*):
                The trial run or the hyperparameter dictionary for hyperparameter search.
            ignore_keys_for_eval (`list[str]`, *optional*)
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions for evaluation during the training.

        Returns:
            [`~trainer_utils.TrainOutput`]: Object containing the global step count, training loss, and metrics.
        """
        if resume_from_checkpoint is False:
            resume_from_checkpoint = None

        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

        args = self.args

        self.is_in_train = True

        # Model re-init
        if self.model_init is not None:
            # Seed must be set before instantiating the model when using model_init.
            enable_full_determinism(args.seed) if args.full_determinism else set_seed(args.seed)
            self.model = self.call_model_init(trial)
            # Reinitializes optimizer and scheduler
            self.optimizer, self.lr_scheduler = None, None
            if self.place_model_on_device:
                self._move_model_to_device(self.model, args.device)
            self.model_wrapped = self.model

        if self.args.use_liger_kernel:
            apply_liger_kernel(self.model, self.args.liger_kernel_config)

        # When fp16/bf16 full eval is enabled, __init__ skips device placement so that
        # evaluation_loop can cast dtype and move in one step. Move the model now for training.
        if (args.fp16_full_eval or args.bf16_full_eval) and not self.is_model_parallel and self.model_init is None:
            self._move_model_to_device(self.model, args.device)

        # Activate gradient checkpointing if needed
        if args.gradient_checkpointing:
            self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=args.gradient_checkpointing_kwargs)

        # If the model uses a tokenizer, it may have a new tokens for fine-tuning purposes.
        if isinstance(self.processing_class, (PreTrainedTokenizerBase, ProcessorMixin)) and hasattr(
            self.model, "config"
        ):
            align_special_tokens(self.model, self.processing_class)

        # Attach NEFTune hooks if necessary
        if self.neftune_noise_alpha is not None:
            self.neftune_hook_handle = activate_neftune(self.model, self.neftune_noise_alpha, self.accelerator)

        # This might change the seed so needs to run first.
        self._hp_search_setup(trial)

        if DebugOption.UNDERFLOW_OVERFLOW in args.debug:
            if args.n_gpu > 1:
                # nn.DataParallel(model) replicates the model, creating new variables and module
                # references registered here no longer work on other gpus, breaking the module
                raise ValueError(
                    "Currently --debug underflow_overflow is not supported under DP. Please use DDP with torchrun"
                )
            else:
                DebugUnderflowOverflow(self.model)

        # Load potential model checkpoint
        if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint:
            resume_from_checkpoint = get_last_checkpoint(args.output_dir)
            if resume_from_checkpoint is None:
                raise ValueError(f"No valid checkpoint found in output directory ({args.output_dir})")

        if resume_from_checkpoint is not None:
            # Load model checkpoint before accelerator.prepare() for regular models,
            # so that buffers and parameters are on the right device after prepare.
            # Deepspeed/FSDP models are loaded after prepare in _prepare_for_training.
            if not is_sagemaker_mp_enabled() and not self.is_deepspeed_enabled and not self.is_fsdp_enabled:
                self._load_from_checkpoint(resume_from_checkpoint)
            state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
            if state.train_batch_size is not None and args.auto_find_batch_size:
                # Only restore the checkpoint's train_batch_size when using auto_find_batch_size,
                self._train_batch_size = state.train_batch_size

        inner_training_loop = find_executable_batch_size(
            self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size
        )
        # Disable progress bars when uploading models during checkpoints to avoid polluting stdout
        ctx = suppress_progress_bars() if args.push_to_hub else contextlib.nullcontext()
        with ctx:
            return inner_training_loop(
                args=args,
                resume_from_checkpoint=resume_from_checkpoint,
                trial=trial,
                ignore_keys_for_eval=ignore_keys_for_eval,
            )

    def _inner_training_loop(
        self,
        batch_size: int | None = None,
        args: TrainingArguments | None = None,
        resume_from_checkpoint: str | None = None,
        trial: "optuna.Trial | dict[str, Any] | None" = None,
        ignore_keys_for_eval: list[str] | None = None,
    ) -> TrainOutput:
        """Run the actual training loop: forward, backward, optimizer step, logging, and checkpointing."""
        # reset everything
        self.accelerator.free_memory()
        if args.auto_find_batch_size:
            self._update_auto_batch_size(batch_size)
        # Data loader and number of training steps
        train_dataloader = self.get_train_dataloader()
        if self.is_fsdp_xla_v2_enabled:
            train_dataloader = tpu_spmd_dataloader(train_dataloader)

        # Setting up training control variables:
        (
            num_train_epochs,
            num_update_steps_per_epoch,
            num_examples,
            num_train_samples,
            total_train_batch_size,
            steps_in_epoch,
            max_steps,
        ) = self.set_initial_training_values(args, train_dataloader)

        epochs_trained, steps_trained_in_current_epoch = self._init_training_state(
            max_steps, num_update_steps_per_epoch, num_train_epochs, resume_from_checkpoint, trial
        )
        model, train_dataloader = self._prepare_for_training(max_steps, train_dataloader, resume_from_checkpoint)

        # Train!
        logger.info("***** Running training *****")
        logger.info(f"  Num examples = {num_examples:,}")
        logger.info(f"  Num Epochs = {num_train_epochs:,}")
        logger.info(f"  Num update steps per epoch = {num_update_steps_per_epoch:,}")
        logger.info(f"  Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}")
        if self.args.per_device_train_batch_size != self._train_batch_size:
            logger.info(f"  Training with DataParallel so batch size has been adjusted to: {self._train_batch_size:,}")
        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size:,}")
        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        logger.info(f"  Total optimization steps = {max_steps:,}")
        logger.info(f"  Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}")

        if resume_from_checkpoint is not None:
            logger.info(
                f"  Resuming training from checkpoint with epoch {epochs_trained} and global step {self.state.global_step}"
            )
            if not self.args.ignore_data_skip:
                logger.info(
                    f"  Fast-forwarding the dataloader past {epochs_trained} epochs and"
                    f" {steps_trained_in_current_epoch} batches to resume from the exact training state."
                )

        start_time = time.time()
        # needed to calculate tokens/s
        self._initial_num_input_tokens_seen = self.state.num_input_tokens_seen
        # Logging state: _tr_loss accumulates on-device between logging steps (avoiding costly .item() syncs
        # on TPUs), then gets drained into _total_loss_scalar at each logging step.
        self._tr_loss = torch.tensor(0.0, device=args.device)
        self._total_loss_scalar = 0.0
        self._globalstep_last_logged = self.state.global_step

        model.zero_grad()

        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)

        if args.eval_on_start:
            self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True)

        for epoch in range(epochs_trained, num_train_epochs):
            self.control = self.callback_handler.on_epoch_begin(self.args, self.state, self.control)
            self._run_epoch(
                model=model,
                epoch=epoch,
                train_dataloader=train_dataloader,
                steps_in_epoch=steps_in_epoch,
                num_update_steps_per_epoch=num_update_steps_per_epoch,
                trial=trial,
                ignore_keys_for_eval=ignore_keys_for_eval,
                start_time=start_time,
                resume_from_checkpoint=resume_from_checkpoint,
                epochs_trained=epochs_trained,
                steps_trained_in_current_epoch=steps_trained_in_current_epoch,
            )
            if self.control.should_training_stop:
                break

        return self._finalize_training(trial, num_train_samples, start_time)

    def _init_training_state(
        self, max_steps, num_update_steps_per_epoch, num_train_epochs, resume_from_checkpoint, trial
    ) -> tuple[int, int]:
        """Initialize TrainerState, optionally restoring from checkpoint. Returns (epochs_trained, steps_trained_in_current_epoch)."""
        self.state = TrainerState(
            stateful_callbacks=[
                cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
            ]
        )
        self.state.is_hyper_param_search = trial is not None
        self.state.train_batch_size = self._train_batch_size
        self.state.compute_steps(self.args, max_steps)

        epochs_trained = 0
        steps_trained_in_current_epoch = 0

        if resume_from_checkpoint is not None and os.path.isfile(
            os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)
        ):
            self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
            compare_trainer_and_checkpoint_args(self.args, self.state)
            self._load_callback_state()
            epochs_trained = int(self.state.global_step // num_update_steps_per_epoch)
            if not self.args.ignore_data_skip:
                steps_trained_in_current_epoch = self.state.global_step % num_update_steps_per_epoch
                steps_trained_in_current_epoch *= self.args.gradient_accumulation_steps

        self.state.init_training_references(self, max_steps, num_train_epochs, trial)

        return epochs_trained, steps_trained_in_current_epoch

    def _prepare_for_training(self, max_steps, train_dataloader, resume_from_checkpoint):
        """Wrap model, create optimizer and scheduler, and run accelerator.prepare. Returns (model, train_dataloader)."""
        delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled

        # Can't delay optimizer creation when using FSDP2: https://github.com/huggingface/accelerate/blob/3f636d626063ffcf9a337c7d3624d61b7d187d59/src/accelerate/accelerator.py#L1404
        is_fsdp2 = self.is_fsdp_enabled and (getattr(self.accelerator.state.fsdp_plugin, "fsdp_version", 1) == 2)
        if is_fsdp2:
            delay_optimizer_creation = False

        # We need to reset the scheduler, as its parameters may be different on subsequent calls
        if self._created_lr_scheduler:
            self.lr_scheduler = None
            self._created_lr_scheduler = False

        if self.is_deepspeed_enabled:
            self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps)

        if not delay_optimizer_creation:
            self.create_optimizer()

        # Pass `self.model_wrapped` so that `_wrap_model` can detect if the model is already
        # wrapped (e.g. in DataParallel) on subsequent `train()` calls and avoid double wrapping.
        model = self._wrap_model(self.model_wrapped)

        # If the model is wrapped, don't use `accelerator.prepare`
        # this is for unhandled cases in accelerate such as FSDP-XLA, SageMaker MP/DP, DataParallel
        use_accelerator_prepare = model is self.model

        # prepare using `accelerator` prepare
        if use_accelerator_prepare:
            if delay_optimizer_creation:
                # TODO: check if we can move this somewhere else
                if self.is_fsdp_enabled and _is_peft_model(self.model):
                    update_fsdp_plugin_peft(self.model, self.accelerator)
                # we only prepare the model as we don't have an optimizer
                model = self.accelerator.prepare(self.model)
                # using the model we prepared to create the optimizer
                self.create_optimizer(model)
                self.optimizer = self.accelerator.prepare(self.optimizer)
            elif self.is_deepspeed_enabled and type(self.lr_scheduler).__name__ == "DummyScheduler":
                model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
                    self.model, self.optimizer, self.lr_scheduler
                )
            else:
                model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
        else:
            self.optimizer = self.accelerator.prepare(self.optimizer)

        # Create scheduler now that the optimizer won't change anymore
        self.create_scheduler(num_training_steps=max_steps)

        # updating self.model_wrapped
        self.model_wrapped = model

        if self.is_fsdp_enabled or self.is_fsdp_xla_enabled:
            # breaking convention for FSDP model
            # TODO: check if this is really needed
            self.model = self.model_wrapped = model

        # backward compatibility
        # TODO: check if we really need this
        if self.is_deepspeed_enabled:
            self.deepspeed = self.model_wrapped

        # Important: at this point:
        # self.model         is the Transformers Model except when we are using FSDP
        # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model),
        # FSDP(Transformers Model), Dynamo Optimized Module(Transformers Model) etc.

        if self.is_fsdp_enabled:
            # Fix `got mixed torch.Tensor and DTensor` error in model.generate() for FSDP2 with LoRA
            if hasattr(self.model, "generate"):
                dist.fsdp.register_fsdp_forward_method(self.model, "generate")

        # since DataLoader was Accelerate prepared w/o a model arg in the same call, we now have to complete the DL wrapping for ALST/UlyssesSP, after model has been prepared
        pc = getattr(self.accelerator, "parallelism_config", None)
        if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled:
            train_dataloader = self.accelerator.deepspeed_ulysses_dl_adapter(train_dataloader, model)

        # load checkpoint
        if resume_from_checkpoint is not None:
            if self.is_deepspeed_enabled:
                deepspeed_load_checkpoint(
                    self.model_wrapped, resume_from_checkpoint, load_module_strict=not _is_peft_model(self.model)
                )
            elif is_sagemaker_mp_enabled() or self.is_fsdp_enabled:
                self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped)

            self._load_optimizer_and_scheduler(resume_from_checkpoint)
            self._load_scaler(resume_from_checkpoint)

        # Update the references for the callback_handler
        for attr in ("model", "optimizer", "lr_scheduler"):
            setattr(self.callback_handler, attr, getattr(self, attr))
        self.callback_handler.train_dataloader = train_dataloader

        return model, train_dataloader

    def _run_epoch(
        self,
        model,
        epoch,
        train_dataloader,
        steps_in_epoch,
        num_update_steps_per_epoch,
        trial,
        ignore_keys_for_eval,
        start_time,
        resume_from_checkpoint,
        epochs_trained,
        steps_trained_in_current_epoch,
    ):
        """Run one full pass over the dataloader."""

        step = -1
        grad_norm = None
        learning_rate = None
        rng_to_sync = False

        # Handle resumption from checkpoint: skip already-trained batches in the resumed epoch
        num_update_steps_trained = 0
        if epoch == epochs_trained and resume_from_checkpoint is not None:
            if steps_trained_in_current_epoch > 0 and not self.args.ignore_data_skip:
                train_dataloader = skip_first_batches(train_dataloader, steps_trained_in_current_epoch)
                step = steps_trained_in_current_epoch - 1
                num_update_steps_trained = steps_trained_in_current_epoch // self.args.gradient_accumulation_steps
                rng_to_sync = True
            elif steps_trained_in_current_epoch == 0:
                self._load_rng_state(resume_from_checkpoint)

        if hasattr(train_dataloader, "set_epoch"):
            train_dataloader.set_epoch(epoch)
        epoch_iterator = iter(train_dataloader)

        # We chunkify the epoch iterator into gradient accumulation steps `n` batches
        remainder = steps_in_epoch % self.args.gradient_accumulation_steps
        if remainder == 0:
            remainder = self.args.gradient_accumulation_steps

        # Outer loop: one iteration per optimizer step. Each iteration prefetches
        # `gradient_accumulation_steps` batches (fewer for the last step if the epoch
        # doesn't divide evenly).
        for update_step in range(num_update_steps_trained, num_update_steps_per_epoch):
            num_batches = (
                self.args.gradient_accumulation_steps if update_step != (num_update_steps_per_epoch - 1) else remainder
            )
            batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches, self.args.device)

            # This is used to correctly scale the loss when the last accumulation step has fewer batches.
            # Not used if `num_items_in_batch` is not None.
            self.current_gradient_accumulation_steps = len(batch_samples)

            # need to sync after if we skipped the batches in `get_batch_samples` for shuffle order reason
            if rng_to_sync:
                self._load_rng_state(resume_from_checkpoint)
                rng_to_sync = False

            # Inner loop: forward + backward for each micro-batch. Gradients are
            # accumulated without syncing until the last micro-batch, then we clip,
            # step the optimizer, and log/save/evaluate.
            for i, inputs in enumerate(batch_samples):
                step += 1
                do_sync_step = (step + 1) % self.args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch
                # Since we perform prefetching, we need to manually set sync_gradients
                self.accelerator.gradient_state._set_sync_gradients(do_sync_step)

                if step % self.args.gradient_accumulation_steps == 0:
                    self.control = self.callback_handler.on_step_begin(self.args, self.state, self.control)

                # We sync the gradients in the following cases: 1. sync_each_batch set to True 2. Using deepspeed 3. when we are at the last batch sample
                if (
                    self.accelerator.gradient_state.plugin_kwargs.get("sync_each_batch", False)
                    or self.accelerator.distributed_type == DistributedType.DEEPSPEED
                    or i == len(batch_samples) - 1
                ):
                    sync_context = contextlib.nullcontext
                else:
                    sync_context = functools.partial(self.accelerator.no_sync, model=model)
                with sync_context():
                    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)

                if (
                    self.args.logging_nan_inf_filter
                    and not is_torch_xla_available()
                    and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
                ):
                    # if loss is nan or inf simply add the average of previous logged losses
                    self._tr_loss += self._tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)
                else:
                    if self._tr_loss.device != tr_loss_step.device:
                        raise ValueError(
                            f"Calculated loss must be on the original device: {self._tr_loss.device} but device in use is {tr_loss_step.device}"
                        )
                    self._tr_loss += tr_loss_step

                self.current_flos += float(self.floating_point_ops(inputs))
                self._track_num_input_tokens(inputs)

                if do_sync_step:
                    grad_norm = None
                    if self.args.max_grad_norm > 0:
                        grad_norm = self._clip_grad_norm(model)
                    grad_norm = self._get_grad_norm(model, grad_norm=grad_norm)

                    self.control = self.callback_handler.on_pre_optimizer_step(self.args, self.state, self.control)
                    self.optimizer.step()
                    self.control = self.callback_handler.on_optimizer_step(self.args, self.state, self.control)

                    # get leaning rate before update
                    learning_rate = self._get_learning_rate()

                    if not self.accelerator.optimizer_step_was_skipped:
                        # Delay optimizer scheduling until metrics are generated
                        if not isinstance(self.lr_scheduler, (torch.optim.lr_scheduler.ReduceLROnPlateau, GreedyLR)):
                            self.lr_scheduler.step()

                    model.zero_grad()
                    self.state.global_step += 1
                    self.state.epoch = epoch + (step + 1) / steps_in_epoch
                    self.control = self.callback_handler.on_step_end(self.args, self.state, self.control)
                    self._maybe_log_save_evaluate(
                        self._tr_loss,
                        grad_norm,
                        model,
                        trial,
                        epoch,
                        ignore_keys_for_eval,
                        start_time,
                        learning_rate=learning_rate,
                    )
                else:
                    self.control = self.callback_handler.on_substep_end(self.args, self.state, self.control)

                if self.control.should_epoch_stop or self.control.should_training_stop:
                    break
            if self.control.should_epoch_stop or self.control.should_training_stop:
                break

        # PyTorch/XLA relies on the dataloader to insert mark_step each iteration.
        # When we break out of the loop early, we flush the pending graph manually.
        if is_torch_xla_available():
            xm.mark_step()

        if step < 0:
            logger.warning(
                "There seems not to be a single sample in your epoch_iterator, stopping training at step"
                f" {self.state.global_step}! This is expected if you're using an IterableDataset and set"
                f" num_steps ({self.state.max_steps}) higher than the number of available samples."
            )
            self.control.should_training_stop = True

        self.control = self.callback_handler.on_epoch_end(self.args, self.state, self.control)
        self._maybe_log_save_evaluate(
            self._tr_loss,
            grad_norm,
            model,
            trial,
            epoch,
            ignore_keys_for_eval,
            start_time,
            learning_rate=learning_rate,
        )

    def _finalize_training(self, trial, num_train_samples, start_time):
        """Finalize training: metrics, best-model loading, cleanup. Returns TrainOutput."""
        logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")

        # add remaining tr_loss
        self._total_loss_scalar += self._tr_loss.item()
        effective_global_step = max(self.state.global_step, 0.001)  # Avoid ZeroDivisionError
        train_loss = self._total_loss_scalar / effective_global_step

        metrics = speed_metrics(
            "train",
            start_time,
            num_samples=num_train_samples,
            num_steps=self.state.max_steps,
        )
        self.store_flos()
        metrics["total_flos"] = self.state.total_flos
        metrics["train_loss"] = train_loss

        self._memory_tracker.stop_and_update_metrics(metrics)
        self.log(metrics)

        if self.args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
            self._load_best_model()

        checkpoints_sorted = sort_checkpoints(
            output_dir=self._get_output_dir(trial), best_model_checkpoint=self.state.best_model_checkpoint
        )

        # Delete the last checkpoint when save_total_limit=1 if it's different from the best checkpoint and process allowed to save.
        if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1:
            for checkpoint in checkpoints_sorted:
                if not os.path.samefile(checkpoint, self.state.best_model_checkpoint):
                    logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
                    shutil.rmtree(checkpoint, ignore_errors=True)

        self.control = self.callback_handler.on_train_end(self.args, self.state, self.control)

        # Wait for the checkpoint to be uploaded.
        self._finish_current_push()

        # After training we make sure to retrieve back the original forward pass method
        # for the embedding layer by removing the forward post hook.
        if self.neftune_noise_alpha is not None:
            deactivate_neftune(self.model, self.neftune_hook_handle, self.accelerator)
        self.is_in_train = False

        return TrainOutput(self.state.global_step, train_loss, metrics)

    def training_step(
        self,
        model: nn.Module,
        inputs: dict[str, torch.Tensor | Any],
        num_items_in_batch: torch.Tensor | int | None = None,
    ) -> torch.Tensor:
        """
        Perform a training step on a batch of inputs.

        Subclass and override to inject custom behavior.

        Args:
            model (`nn.Module`):
                The model to train.
            inputs (`dict[str, torch.Tensor | Any]`):
                The inputs and targets of the model.

                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.

        Return:
            `torch.Tensor`: The tensor with training loss on this batch.
        """
        # Prepare buffers for context parallelism

        cp_context, inputs = self._prepare_context_parallel_inputs(model, inputs)

        # Context manager is no-op if CP isn't enabled
        with cp_context():
            model.train()
            if hasattr(self.optimizer, "train") and callable(self.optimizer.train):
                self.optimizer.train()

            inputs = self._prepare_inputs(inputs)
            if is_sagemaker_mp_enabled():
                loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
                return loss_mb.reduce_mean().detach().to(self.args.device)

            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)

            del inputs
            if (
                self.args.torch_empty_cache_steps is not None
                and self.state.global_step % self.args.torch_empty_cache_steps == 0
            ):
                clear_device_cache()

            kwargs = {}

            # For LOMO optimizers you need to explicitly use the learning rate
            if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
                kwargs["learning_rate"] = self._get_learning_rate()

            if self.args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training

            # Finally we need to normalize the loss for reporting if GA loss bug is not fixed during compute loss
            if (not self.model_accepts_loss_kwargs or num_items_in_batch is None) and self.compute_loss_func is None:
                # If the model does not accept loss kwargs, we need to normalize the loss by the number of gradient accumulation steps
                loss = loss / self.current_gradient_accumulation_steps

            # Turning off loss scaling w.r.t. gradient accumulation when DeepSpeed is enabled
            # https://github.com/huggingface/transformers/pull/35808
            if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
                kwargs["scale_wrt_gas"] = False

            self.accelerator.backward(loss, **kwargs)

            return loss.detach()

    def compute_loss(
        self,
        model: nn.Module,
        inputs: dict[str, torch.Tensor | Any],
        return_outputs: bool = False,
        num_items_in_batch: torch.Tensor | int | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, Any]:
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.

        Args:
            model (`nn.Module`):
                The model to compute the loss for.
            inputs (`dict[str, torch.Tensor | Any]`):
                The input data for the model.
            return_outputs (`bool`, *optional*, defaults to `False`):
                Whether to return the model outputs along with the loss.
            num_items_in_batch (Optional[torch.Tensor], *optional*):
                The number of items in the batch. If not passed, the loss is computed
                using the default batch size reduction logic.

        Returns:
            The loss of the model along with its output if return_outputs was set to True

        Subclass and override for custom behavior. If you are not using `num_items_in_batch` when computing your loss,
        make sure to overwrite `self.model_accepts_loss_kwargs` to `False`. Otherwise, the loss calculation might be slightly inaccurate when performing gradient accumulation.
        """
        pc = getattr(self.accelerator, "parallelism_config", None)
        if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled and self.model.training:
            return deepspeed_sp_compute_loss(self.accelerator, model, inputs, return_outputs, pc)

        if (self.label_smoother is not None or self.compute_loss_func is not None) and "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            labels = None
        if self.model_accepts_loss_kwargs:
            kwargs = {}
            if num_items_in_batch is not None:
                kwargs["num_items_in_batch"] = num_items_in_batch
            inputs = {**inputs, **kwargs}
        outputs = model(**inputs)

        # User-defined compute_loss function
        if self.compute_loss_func is not None:
            if labels is None:
                logger.warning(
                    "Trainer: `compute_loss_func` is defined but `labels=None`. "
                    "Your custom loss function will still be called with labels=None. "
                )
            loss = self.compute_loss_func(
                outputs,
                labels,
                num_items_in_batch=num_items_in_batch,
            )
        # Default HF loss handling (label smoothing) if no custom loss function
        elif labels is not None:
            unwrapped_model = self.accelerator.unwrap_model(model)
            model_name = (
                unwrapped_model.base_model.model._get_name()
                if _is_peft_model(unwrapped_model)
                else unwrapped_model._get_name()
            )
            if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
                loss = self.label_smoother(outputs, labels, shift_labels=True)
            else:
                loss = self.label_smoother(outputs, labels)
        else:
            if isinstance(outputs, dict) and "loss" not in outputs:
                raise ValueError(
                    "The model did not return a loss from the inputs, only the following keys: "
                    f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
                )
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

        if (
            self.args.average_tokens_across_devices
            and (self.model_accepts_loss_kwargs or self.compute_loss_func)
            and num_items_in_batch is not None
        ):
            loss *= self.accelerator.num_processes if self.args.n_gpu <= 1 else self.args.n_gpu

        return (loss, outputs) if return_outputs else loss

    def compute_loss_context_manager(self) -> contextlib.ExitStack:
        """
        A helper wrapper to group together context managers.
        """
        ctx_stack = contextlib.ExitStack()

        autocast_ctx = self.autocast_smart_context_manager()
        if not isinstance(autocast_ctx, contextlib.nullcontext):
            ctx_stack.enter_context(autocast_ctx)

        return ctx_stack

    def autocast_smart_context_manager(self, cache_enabled: bool | None = True) -> contextlib.AbstractContextManager:
        """
        A helper wrapper that creates an appropriate context manager for `autocast` while feeding it the desired
        arguments, depending on the situation. We rely on accelerate for autocast, hence we do nothing here.
        """
        return contextlib.nullcontext()

    def _maybe_log_save_evaluate(
        self,
        tr_loss: torch.Tensor,
        grad_norm: torch.Tensor | float | None,
        model: nn.Module,
        trial: "optuna.Trial | dict[str, Any] | None",
        epoch: float,
        ignore_keys_for_eval: list[str] | None,
        start_time: float,
        learning_rate: float | None = None,
    ) -> None:
        """Log metrics, run evaluation, and save checkpoints if the current training state requires it."""
        if self.control.should_log and self.state.global_step > self._globalstep_last_logged:
            if is_torch_xla_available():
                xm.mark_step()

            logs: dict[str, float] = {}

            # all_gather + mean() to get average loss over all processes
            tr_loss_scalar = nested_gather(tr_loss, self.args.parallel_mode).mean().item()

            # reset tr_loss to zero
            tr_loss -= tr_loss

            logs["loss"] = tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged)
            if grad_norm is not None:
                logs["grad_norm"] = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm
            if learning_rate is not None:
                logs["learning_rate"] = learning_rate
            else:
                logs["learning_rate"] = self._get_learning_rate()

            self._total_loss_scalar += tr_loss_scalar
            self._globalstep_last_logged = self.state.global_step
            self.store_flos()

            self.log(logs, start_time)

        metrics = None
        if self.control.should_evaluate:
            metrics = self._evaluate(trial, ignore_keys_for_eval)
            is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial)

            if self.args.save_strategy == SaveStrategy.BEST:
                self.control.should_save = is_new_best_metric

        if self.control.should_save:
            self._save_checkpoint(model, trial)
            self.control = self.callback_handler.on_save(self.args, self.state, self.control)

    # ---- Training Utilites ----
    def get_batch_samples(
        self, epoch_iterator: Iterator, num_batches: int, device: torch.device
    ) -> tuple[list, torch.Tensor | int | None]:
        """
        Collects a specified number of batches from the epoch iterator and optionally counts the number of items in the batches to properly scale the loss.
        """
        batch_samples = []

        for _ in range(num_batches):
            try:
                batch_samples.append(next(epoch_iterator))
            except StopIteration:
                break

        num_items_in_batch = self._get_num_items_in_batch(batch_samples, device)
        return batch_samples, num_items_in_batch

    def _get_num_items_in_batch(self, batch_samples: list, device: torch.device) -> torch.Tensor | int | None:
        """
        Counts the number of items in the batches to properly scale the loss.
        Args:
            batch_samples (`list`): List of batches
            device (`torch.device`): The device on which the number of items in the batch should be.
        Returns:
            None if the number of items in the batch doesn't need to be computed else the number of items in the batch
        """
        num_items_in_batch = None
        count_num_items_in_batch = (
            len(batch_samples) > 0
            and "labels" in batch_samples[0]
            and (
                # num_items_in_batch is passed to model forward
                # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3757
                self.model_accepts_loss_kwargs
                # num_items_in_batch is passed to compute_loss_func
                # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3773
                or self.compute_loss_func is not None
                # num_items_in_batch is also verified if (self.model_accepts_loss_kwargs or self.compute_loss_func)
                # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3790
            )
        )
        if count_num_items_in_batch:
            # For now we don't support object detection
            try:
                num_items_in_batch = sum((batch["labels"].ne(-100)).sum() for batch in batch_samples)
            except (TypeError, AttributeError):
                pass

        if num_items_in_batch is not None:
            if self.args.average_tokens_across_devices:
                if self.args.world_size > 1:
                    num_items_in_batch = self.accelerator.gather(num_items_in_batch.to(device)).sum()
            elif self.args.n_gpu > 1:
                # In DP case, if we don't average, we need to divide by the number of gpu. This is the simplest approximation.
                # Otherwise, we would have to scatter labels and calculate num_items_in_batch for each gpu.
                num_items_in_batch = num_items_in_batch // self.args.n_gpu

            if torch.is_tensor(num_items_in_batch):
                num_items_in_batch = num_items_in_batch.to(device)

                if self.args.n_gpu > 1 and num_items_in_batch.dim() == 0:
                    # In the DataParallel case, convert the scalar tensor into a 2-dim tensor with the same value repeated
                    num_items_in_batch = num_items_in_batch.unsqueeze(0).expand(self.args.n_gpu, -1)
                # Divide by number of devices with the same batch
                if pc := getattr(self.accelerator, "parallelism_config", None):
                    num_items_in_batch = num_items_in_batch // pc.non_data_parallel_size

        return num_items_in_batch

    def _prepare_input(self, data: torch.Tensor | Any) -> torch.Tensor | Any:
        """
        Prepares one `data` before feeding it to the model, be it a tensor or a nested list/dictionary of tensors.
        """
        if isinstance(data, Mapping):
            return type(data)({k: self._prepare_input(v) for k, v in data.items()})
        elif isinstance(data, (tuple, list)):
            return type(data)(self._prepare_input(v) for v in data)
        elif isinstance(data, torch.Tensor):
            kwargs = {"device": self.args.device}
            if self.is_deepspeed_enabled and (torch.is_floating_point(data) or torch.is_complex(data)):
                # NLP models inputs are int/uint and those get adjusted to the right dtype of the
                # embedding. Other models such as wav2vec2's inputs are already float and thus
                # may need special handling to match the dtypes of the model
                kwargs.update({"dtype": self.accelerator.state.deepspeed_plugin.hf_ds_config.dtype()})
            return data.to(**kwargs)
        return data

    def _prepare_inputs(self, inputs: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]:
        """
        Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and
        handling potential state.
        """
        inputs = self._prepare_input(inputs)
        if len(inputs) == 0:
            raise ValueError(
                "The batch received was empty, your model won't be able to train on it. Double-check that your "
                f"training dataset contains keys expected by the model: {','.join(self._signature_columns)}."
            )

        return inputs

    def _prepare_context_parallel_inputs(
        self, model: nn.Module, inputs: dict[str, torch.Tensor | Any]
    ) -> tuple[Callable, dict[str, torch.Tensor | Any]]:
        """
        Prepare inputs for context parallelism by setting up buffers and validation.

        Args:
            model: The model being trained
            inputs: Input tensors to prepare

        Returns:
            tuple: (context_manager, prepared_inputs) where context_manager is either
                   the context parallelism wrapper or a no-op context
        """
        if (
            getattr(self.accelerator, "parallelism_config", None) is not None
            and self.accelerator.parallelism_config.cp_enabled
        ):
            if self.accelerator.parallelism_config.cp_backend == "torch":
                if hasattr(model, "config"):
                    if model.config._attn_implementation != "sdpa":
                        raise ValueError(
                            f"Context parallelism is supported only with SDPA attention, you are using {model.config._attn_implementation}."
                        )

                if "shift_labels" not in inputs:
                    logger.warning_once("Shift labels not found in the inputs, shifting manually")
                    if "labels" in inputs:
                        _ignore_index = -100
                        labels = nn.functional.pad(inputs["labels"], (0, 1), value=_ignore_index)
                        inputs["shift_labels"] = labels[:, 1:].contiguous()

            # note: we don't do anything for accelerator.parallelism_config.sp_backend == "deepspeed" since:
            # - accelerator.parallelism_config performs the `model.config._attn_implementation` checks already and it supports more than `dspa`
            # - UlyssesSPDataLoaderAdapter called from Accelerate performs the `shift_label` creation - must not interfere
            # - position_ids generation should be done by HF Trainer if it wasn't done by the user

            if "position_ids" not in inputs:
                logger.warning_once("Position IDs not found in the inputs, generating manually")
                inputs["position_ids"] = torch.arange(
                    inputs["input_ids"].size(1), device=inputs["input_ids"].device
                ).expand(inputs["input_ids"].size(0), -1)

            buffers = []
            buffer_seq_dims = []

            if "input_ids" in inputs:
                buffers.append(inputs["input_ids"])
                buffer_seq_dims.append(1)  # Sequence dimension
            if "labels" in inputs:
                buffers.append(inputs["labels"])
                buffer_seq_dims.append(1)
            if "shift_labels" in inputs:
                buffers.append(inputs["shift_labels"])
                buffer_seq_dims.append(1)
            # Add attention_mask to buffers for context parallel splitting (only if causal)
            if "attention_mask" in inputs:
                # Only validate causal mask once for performance
                if not getattr(self, "_attn_mask_causal_checked", False):
                    # Context parallel currently doesn't support other masks than causal
                    # Accelerate applies hooks to replace mask with is_causal arg in SDPA
                    # Check if the mask is really causal and if not throw an error
                    attention_mask = inputs["attention_mask"]
                    if not is_attention_mask_causal(attention_mask):
                        raise ValueError(
                            "Context parallelism only supports causal attention masks. "
                            "The provided attention_mask is not causal. "
                            "Please ensure your data uses causal masking (lower triangular) "
                            "or remove the attention_mask to use the model's default causal masking."
                        )
                    self._attn_mask_causal_checked = True
                if self._attn_mask_causal_checked:
                    # Add to buffers only after validation (or if validation already passed)
                    attention_mask = inputs["attention_mask"]
                    if attention_mask.dim() == 2:
                        buffers.append(attention_mask)
                        buffer_seq_dims.append(1)
                    else:
                        # Other dimensionality; keep as-is without sharding to avoid incorrect splits
                        pass
            # Include position_ids in context parallelism splitting
            if "position_ids" in inputs and inputs["position_ids"] is not None:
                buffers.append(inputs["position_ids"])
                buffer_seq_dims.append(1)

            return partial(
                self.accelerator.maybe_context_parallel,
                buffers=buffers,
                buffer_seq_dims=buffer_seq_dims,
                no_restore_buffers=set(buffers),
            ), inputs

        return contextlib.nullcontext, inputs

    def set_initial_training_values(
        self, args: TrainingArguments, dataloader: DataLoader
    ) -> tuple[int, int, int, int, int, int | None, int]:
        """
        Calculates and returns the following values:
        - `num_train_epochs`
        - `num_update_steps_per_epoch`
        - `num_examples`
        - `num_train_samples`
        - `total_train_batch_size`
        - `steps_in_epoch` (total batches per epoch)
        - `max_steps`
        """
        # Case 1: we rely on `args.max_steps` first
        max_steps = args.max_steps
        # If max_steps is negative, we use the number of epochs to determine the number of total steps later
        epoch_based = max_steps < 0
        len_dataloader = len(dataloader) if has_length(dataloader) else None
        total_train_batch_size = self.get_total_train_batch_size(args)

        # Account for Sequence Parallelism (SP) dataloader adapter's effect
        sp_size = self.get_sp_size()
        if sp_size > 1 and len_dataloader is not None:
            len_dataloader = len_dataloader * sp_size

        # Case 2: We have a dataloader length and can extrapolate
        if len_dataloader is not None:
            num_update_steps_per_epoch = max(
                len_dataloader // args.gradient_accumulation_steps
                + int(len_dataloader % args.gradient_accumulation_steps > 0),
                1,
            )
            # Case 3: We have a length but are using epochs, we can extrapolate the number of steps
            if epoch_based:
                max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
        # Now we figure out `num_examples`, `num_train_epochs`, and `train_samples`
        if len_dataloader:
            num_examples = self.num_examples(dataloader)
            if args.max_steps > 0:
                num_train_epochs = max_steps // num_update_steps_per_epoch + int(
                    max_steps % num_update_steps_per_epoch > 0
                )
                # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's
                # the best we can do.
                num_train_samples = max_steps * total_train_batch_size
            else:
                num_train_epochs = math.ceil(args.num_train_epochs)
                num_train_samples = self.num_examples(dataloader) * args.num_train_epochs
        elif args.max_steps > 0:  # Rely on max_steps when dataloader does not have a working size
            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
            num_train_epochs = sys.maxsize
            num_update_steps_per_epoch = max_steps
            num_examples = total_train_batch_size * args.max_steps
            num_train_samples = args.max_steps * total_train_batch_size
        else:
            raise ValueError(
                "args.max_steps must be set to a positive value if dataloader does not have a length, was"
                f" {args.max_steps}"
            )
        steps_in_epoch = len_dataloader if len_dataloader is not None else max_steps * args.gradient_accumulation_steps
        return (
            num_train_epochs,
            num_update_steps_per_epoch,
            num_examples,
            num_train_samples,
            total_train_batch_size,
            steps_in_epoch,
            max_steps,
        )

    def get_total_train_batch_size(self, args: TrainingArguments) -> int:
        """Calculates total batch size (micro_batch * grad_accum * dp_world_size).

        Accounts for all parallelism dimensions: TP, CP, and SP.

        Formula: dp_world_size = world_size // (tp_size * cp_size * sp_size)

        Where:
        - TP (Tensor Parallelism): Model layers split across GPUs
        - CP (Context Parallelism): Sequences split using Ring Attention (FSDP2)
        - SP (Sequence Parallelism): Sequences split using ALST/Ulysses (DeepSpeed)

        All dimensions are separate and multiplicative: world_size = dp_size * tp_size * cp_size * sp_size
        """

        dp_world_size = args.world_size // self.get_tp_size() // self.get_cp_size() // self.get_sp_size()
        return self._train_batch_size * args.gradient_accumulation_steps * dp_world_size

    def get_sp_size(self) -> int:
        """Get the sequence parallel size"""
        if getattr(self.accelerator, "parallelism_config", None) is None:
            return 1
        else:
            pc = self.accelerator.parallelism_config
            return pc.sp_size

    def get_cp_size(self) -> int:
        """Get the context parallel size"""
        if getattr(self.accelerator, "parallelism_config", None) is None:
            return 1
        else:
            pc = self.accelerator.parallelism_config
            return pc.cp_size

    def get_tp_size(self) -> int:
        """Get the tensor parallel size from either the model or DeepSpeed config."""

        # 1. Check model.tp_size first
        if (model_tp := getattr(self.model, "_tp_size", None)) is not None:
            return model_tp

        # 2. Fall back to DeepSpeed config if enabled
        if self.is_deepspeed_enabled and (deepspeed_config := getattr(self.args, "hf_deepspeed_config", None)):
            return deepspeed_config.config.get("tensor_parallel", {}).get("autotp_size", 1)

        # 3. Default fallback
        return 1

    def _wrap_model(self, model: nn.Module, training: bool = True, dataloader: DataLoader | None = None) -> nn.Module:
        """Wrap `model` for distributed training if needed (DDP, FSDP, SageMaker, etc.)."""
        # train/eval could be run multiple-times - if already wrapped, don't re-wrap it again
        if self.accelerator.unwrap_model(model, keep_torch_compile=False) is not model:
            return model

        if is_sagemaker_mp_enabled():
            # Wrapping the base model twice in a DistributedModel will raise an error.
            if isinstance(model, smp.model.DistributedModel):
                return model
            return smp.DistributedModel(model, backward_passes_per_step=self.args.gradient_accumulation_steps)

        # Multi-gpu training, quantized models do not support DP
        if (
            self.args.n_gpu > 1
            and not getattr(model, "is_loaded_in_8bit", False)
            and not getattr(model, "is_loaded_in_4bit", False)
        ):
            model = nn.DataParallel(model)

        # Note: in torch.distributed mode, there's no point in wrapping the model
        # inside a DistributedDataParallel as we'll be under `no_grad` anyways.
        if not training:
            return model

        # Distributed training using PyTorch FSDP
        if self.is_fsdp_xla_enabled:
            model = wrap_model_xla_fsdp(model, self.args, self.is_fsdp_xla_v2_enabled)
        elif is_sagemaker_dp_enabled():
            model = nn.parallel.DistributedDataParallel(
                model, device_ids=[int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"))]
            )
        return model

    def _update_auto_batch_size(self, batch_size):
        """Free memory, reset model wrapping, and update DeepSpeed config for the new batch size when using `auto_find_batch_size`"""
        # `_train_batch_size` value might have changed to `auto_find_batch_size`
        self._train_batch_size = batch_size
        # frees the wrapped model and resets it back to the unwrapped base model
        release_memory(self.model_wrapped)

        if self.is_fsdp_enabled:
            # Remove FSDP wrapping from sub-models because self.model points to the wrapped model in FSDP case
            self.model = unwrap_model(self.model, recursive=True)

        self.model_wrapped = self.model

        # Check for DeepSpeed *after* the initial pass and modify the config
        if self.is_deepspeed_enabled:
            # Temporarily unset `self.args.train_batch_size`
            original_bs = self.args.per_device_train_batch_size
            self.args.per_device_train_batch_size = self._train_batch_size // max(1, self.args.n_gpu)
            propagate_args_to_deepspeed(self.accelerator, self.args, auto_find_batch_size=True)
            self.args.per_device_train_batch_size = original_bs

    def _track_num_input_tokens(self, inputs):
        """Count input tokens seen (all or non-padding) and update state."""
        if self.args.include_num_input_tokens_seen == "no":
            return
        main_input_name = getattr(self.model, "main_input_name", "input_ids")
        if main_input_name not in inputs:
            logger.warning(
                "Tried to track the number of tokens seen, however the current model is "
                "not configured properly to know what item is the input. To fix this, add "
                "a `main_input_name` attribute to the model class you are using."
            )
            return

        if self.args.include_num_input_tokens_seen == "non_padding":
            if "attention_mask" in inputs:
                input_tokens = inputs["attention_mask"].sum()
            elif (
                self.processing_class is not None
                and hasattr(self.processing_class, "pad_token_id")
                and self.processing_class.pad_token_id is not None
            ):
                input_tokens = (inputs[main_input_name] != self.processing_class.pad_token_id).sum()
            else:
                logger.warning(
                    "Could not determine method to count non-padding tokens, falling back to counting all tokens."
                )
                input_tokens = inputs[main_input_name].numel()
        else:
            input_tokens = inputs[main_input_name].numel()

        input_tokens = torch.as_tensor(input_tokens, device=self.args.device, dtype=torch.int64)
        self.state.num_input_tokens_seen += self.accelerator.gather(input_tokens).sum().item()

    def _clip_grad_norm(self, model):
        """Clip gradients to max_grad_norm. Returns the pre-clip gradient norm."""
        if is_sagemaker_mp_enabled() and self.args.fp16:
            return self.optimizer.clip_master_grads(self.args.max_grad_norm)
        return self.accelerator.clip_grad_norm_(model.parameters(), self.args.max_grad_norm)

    def _get_grad_norm(self, model, grad_norm=None):
        """Return the gradient norm as a Python float."""
        if grad_norm is None:
            # Compute norm without clipping (inf means no actual clipping happens)
            grad_norm = self.accelerator.clip_grad_norm_(model.parameters(), float("inf"))

        if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
            if hasattr(grad_norm, "item"):
                grad_norm = grad_norm.item()
        return grad_norm

    # ---- Evaluation & Prediction ----

    def evaluate(
        self,
        eval_dataset: Dataset | dict[str, Dataset] | None = None,
        ignore_keys: list[str] | None = None,
        metric_key_prefix: str = "eval",
    ) -> dict[str, float]:
        """
        Run evaluation and returns metrics.

        The calling script will be responsible for providing a method to compute metrics, as they are task-dependent
        (pass it to the init `compute_metrics` argument).

        You can also subclass and override this method to inject custom behavior.

        Args:
            eval_dataset (`Dataset` | dict[str, `Dataset`], *optional*):
                Pass a dataset if you wish to override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns
                not accepted by the `model.forward()` method are automatically removed. If it is a dictionary, it will
                evaluate on each dataset, prepending the dictionary key to the metric name. Datasets must implement the
                `__len__` method.

                <Tip>

                If you pass a dictionary with names of datasets as keys and datasets as values, evaluate will run
                separate evaluations on each dataset. This can be useful to monitor how training affects other
                datasets or simply to get a more fine-grained evaluation.
                When used with `load_best_model_at_end`, make sure `metric_for_best_model` references exactly one
                of the datasets. If you, for example, pass in `{"data1": data1, "data2": data2}` for two datasets
                `data1` and `data2`, you could specify `metric_for_best_model="eval_data1_loss"` for using the
                loss on `data1` and `metric_for_best_model="eval_data2_loss"` for the loss on `data2`.

                </Tip>

            ignore_keys (`list[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
            metric_key_prefix (`str`, *optional*, defaults to `"eval"`):
                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
                "eval_bleu" if the prefix is "eval" (default)

        Returns:
            A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The
            dictionary also contains the epoch number which comes from the training state.
        """
        # handle multiple eval datasets
        override = eval_dataset is not None
        eval_dataset = eval_dataset if override else self.eval_dataset
        if isinstance(eval_dataset, dict):
            metrics = {}
            for eval_dataset_name, _eval_dataset in eval_dataset.items():
                dataset_metrics = self.evaluate(
                    eval_dataset=_eval_dataset if override else eval_dataset_name,
                    ignore_keys=ignore_keys,
                    metric_key_prefix=f"{metric_key_prefix}_{eval_dataset_name}",
                )
                metrics.update(dataset_metrics)
            return metrics

        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

        eval_dataloader = self.get_eval_dataloader(eval_dataset)
        if self.is_fsdp_xla_v2_enabled:
            eval_dataloader = tpu_spmd_dataloader(eval_dataloader)

        start_time = time.time()

        output = self.evaluation_loop(
            eval_dataloader,
            description="Evaluation",
            # No point gathering the predictions if there are no metrics, otherwise we defer to
            # self.args.prediction_loss_only
            prediction_loss_only=True if self.compute_metrics is None else None,
            ignore_keys=ignore_keys,
            metric_key_prefix=metric_key_prefix,
        )

        total_batch_size = self.args.eval_batch_size * self.args.world_size
        if f"{metric_key_prefix}_model_preparation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"]
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )

        self.log(output.metrics)

        if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
            xm.master_print(met.metrics_report())

        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)

        self._memory_tracker.stop_and_update_metrics(output.metrics)

        return output.metrics

    def evaluation_loop(
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: bool | None = None,
        ignore_keys: list[str] | None = None,
        metric_key_prefix: str = "eval",
    ) -> EvalLoopOutput:
        """
        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.

        Works both with or without labels.
        """
        args = self.args

        prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only

        # if eval is called w/o train, handle model prep here
        if self.is_deepspeed_enabled and self.deepspeed is None:
            _, _ = deepspeed_init(self, num_training_steps=0, inference=True)

        model = self._wrap_model(self.model, training=False)

        if len(self.accelerator._models) == 0 and model is self.model:
            start_time = time.time()
            model = (
                self.accelerator.prepare(model)
                if self.is_deepspeed_enabled or (self.is_fsdp_enabled and not self.args.torch_compile)
                else self.accelerator.prepare_model(model, evaluation_mode=True)
            )
            self.model_preparation_time = round(time.time() - start_time, 4)

            if self.is_fsdp_enabled:
                self.model = model

            # for the rest of this function `model` is the outside model, whether it was wrapped or not
            if model is not self.model:
                self.model_wrapped = model

            # backward compatibility
            if self.is_deepspeed_enabled:
                self.deepspeed = self.model_wrapped

        # if full fp16 or bf16 eval is wanted and this ``evaluation`` or ``predict`` isn't called
        # while ``train`` is running, cast it to the right dtype first and then put on device
        if not self.is_in_train:
            if args.fp16_full_eval:
                model = model.to(dtype=torch.float16, device=args.device)
            elif args.bf16_full_eval:
                model = model.to(dtype=torch.bfloat16, device=args.device)

        batch_size = self.args.eval_batch_size

        logger.info(f"\n***** Running {description} *****")
        if has_length(dataloader):
            logger.info(f"  Num examples = {self.num_examples(dataloader)}")
        else:
            logger.info("  Num examples: Unknown")
        logger.info(f"  Batch size = {batch_size}")

        if hasattr(model, "eval") and callable(model.eval):
            model.eval()
        if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval):
            self.optimizer.eval()

        self.callback_handler.eval_dataloader = dataloader
        # Do this before wrapping.
        eval_dataset = getattr(dataloader, "dataset", None)

        # Initialize containers
        all_losses = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
        all_preds = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
        all_labels = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
        all_inputs = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)

        metrics = None
        eval_set_kwargs = {}

        # Will be useful when we have an iterable dataset so don't know its length.
        observed_num_examples = 0

        # Main evaluation loop
        for step, inputs in enumerate(dataloader):
            # Update the observed num examples
            observed_batch_size = find_batch_size(inputs)
            if observed_batch_size is not None:
                observed_num_examples += observed_batch_size
                # For batch samplers, batch_size is not known by the dataloader in advance.
                if batch_size is None:
                    batch_size = observed_batch_size

            # Prediction step
            losses, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
            main_input_name = getattr(self.model, "main_input_name", "input_ids")
            inputs_decode = (
                self._prepare_input(inputs[main_input_name]) if "inputs" in args.include_for_metrics else None
            )

            if is_torch_xla_available():
                xm.mark_step()

            # Update containers
            if losses is not None:
                losses = self.gather_function(losses.repeat(batch_size))
                all_losses.add(losses)
            if inputs_decode is not None:
                inputs_decode = self.accelerator.pad_across_processes(inputs_decode, dim=1, pad_index=-100)
                inputs_decode = self.gather_function(inputs_decode)
                if not self.args.batch_eval_metrics or description == "Prediction":
                    all_inputs.add(inputs_decode)
            if labels is not None:
                # Pad labels here, preparing for preprocess_logits_for_metrics in next logits block.
                labels = self.accelerator.pad_across_processes(labels, dim=1, pad_index=-100)
            if logits is not None:
                logits = self.accelerator.pad_across_processes(logits, dim=1, pad_index=-100)
                if self.preprocess_logits_for_metrics is not None:
                    logits = self.preprocess_logits_for_metrics(logits, labels)
                logits = self.gather_function(logits)
                if not self.args.batch_eval_metrics or description == "Prediction":
                    all_preds.add(logits)
            if labels is not None:
                labels = self.gather_function(labels)
                if not self.args.batch_eval_metrics or description == "Prediction":
                    all_labels.add(labels)

            self.control = self.callback_handler.on_prediction_step(args, self.state, self.control)

            if self.args.batch_eval_metrics:
                if self.compute_metrics is not None and logits is not None and labels is not None:
                    is_last_step = self.accelerator.gradient_state.end_of_dataloader
                    batch_kwargs = {}
                    batch_kwargs["losses"] = losses if "loss" in args.include_for_metrics else None
                    batch_kwargs["inputs"] = inputs if "inputs" in args.include_for_metrics else None
                    metrics = self.compute_metrics(
                        EvalPrediction(predictions=logits, label_ids=labels, **batch_kwargs),
                        compute_result=is_last_step,
                    )

                del losses, logits, labels, inputs
                torch.cuda.empty_cache()

            # Gather all tensors and put them back on the CPU if we have done enough accumulation steps.
            elif args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0:
                all_losses.to_cpu_and_numpy()
                all_preds.to_cpu_and_numpy()
                all_labels.to_cpu_and_numpy()
                all_inputs.to_cpu_and_numpy()

                del losses, logits, labels, inputs
                torch.cuda.empty_cache()

        # After all calls to `.gather_function`, reset to `gather_for_metrics`:
        self.gather_function = self.accelerator.gather_for_metrics

        # Gather all remaining tensors and put them back on the CPU
        all_losses = all_losses.get_arrays()
        all_preds = all_preds.get_arrays()
        all_labels = all_labels.get_arrays()
        all_inputs = all_inputs.get_arrays()

        # Number of samples
        if has_length(eval_dataset):
            num_samples = len(eval_dataset)
        # The instance check is weird and does not actually check for the type, but whether the dataset has the right
        # methods. Therefore we need to make sure it also has the attribute.
        elif isinstance(eval_dataset, IterableDatasetShard) and getattr(eval_dataset, "num_examples", 0) > 0:
            num_samples = eval_dataset.num_examples
        else:
            if has_length(dataloader):
                num_samples = self.num_examples(dataloader)
            else:  # both len(dataloader.dataset) and len(dataloader) fail
                num_samples = observed_num_examples
        if num_samples == 0 and observed_num_examples > 0:
            num_samples = observed_num_examples

        # Metrics!
        if (
            self.compute_metrics is not None
            and all_preds is not None
            and all_labels is not None
            and not self.args.batch_eval_metrics
        ):
            eval_set_kwargs["losses"] = all_losses if "loss" in args.include_for_metrics else None
            eval_set_kwargs["inputs"] = all_inputs if "inputs" in args.include_for_metrics else None
            metrics = self.compute_metrics(
                EvalPrediction(predictions=all_preds, label_ids=all_labels, **eval_set_kwargs)
            )
        elif metrics is None:
            metrics = {}

        # To be JSON-serializable, we need to remove numpy types or zero-d tensors
        metrics = denumpify_detensorize(metrics)

        if isinstance(all_losses, list) and all_losses:
            metrics[f"{metric_key_prefix}_loss"] = np.concatenate(all_losses).mean().item()
        elif isinstance(all_losses, np.ndarray):
            metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()
        if hasattr(self, "model_preparation_time"):
            metrics[f"{metric_key_prefix}_model_preparation_time"] = self.model_preparation_time

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

        return EvalLoopOutput(predictions=all_preds, label_ids=all_labels, metrics=metrics, num_samples=num_samples)

    def predict(
        self, test_dataset: Dataset, ignore_keys: list[str] | None = None, metric_key_prefix: str = "test"
    ) -> PredictionOutput:
        """
        Run prediction and returns predictions and potential metrics.

        Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method
        will also return metrics, like in `evaluate()`.

        Args:
            test_dataset (`Dataset`):
                Dataset to run the predictions on. If it is an `datasets.Dataset`, columns not accepted by the
                `model.forward()` method are automatically removed. Has to implement the method `__len__`
            ignore_keys (`list[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
            metric_key_prefix (`str`, *optional*, defaults to `"test"`):
                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
                "test_bleu" if the prefix is "test" (default)

        <Tip>

        If your predictions or labels have different sequence length (for instance because you're doing dynamic padding
        in a token classification task) the predictions will be padded (on the right) to allow for concatenation into
        one array. The padding index is -100.

        </Tip>

        Returns: *NamedTuple* A namedtuple with the following keys:

            - predictions (`np.ndarray`): The predictions on `test_dataset`.
            - label_ids (`np.ndarray`, *optional*): The labels (if the dataset contained some).
            - metrics (`dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained
              labels).
        """
        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

        test_dataloader = self.get_test_dataloader(test_dataset)
        start_time = time.time()

        output = self.evaluation_loop(
            test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
        )
        total_batch_size = self.args.eval_batch_size * self.args.world_size
        if f"{metric_key_prefix}_model_preparation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"]
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )

        self.control = self.callback_handler.on_predict(self.args, self.state, self.control, output.metrics)
        self._memory_tracker.stop_and_update_metrics(output.metrics)

        return PredictionOutput(predictions=output.predictions, label_ids=output.label_ids, metrics=output.metrics)

    def prediction_step(
        self,
        model: nn.Module,
        inputs: dict[str, torch.Tensor | Any],
        prediction_loss_only: bool,
        ignore_keys: list[str] | None = None,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
        """
        Perform an evaluation step on `model` using `inputs`.

        Subclass and override to inject custom behavior.

        Args:
            model (`nn.Module`):
                The model to evaluate.
            inputs (`dict[str, torch.Tensor | Any]`):
                The inputs and targets of the model.

                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.
            prediction_loss_only (`bool`):
                Whether or not to return the loss only.
            ignore_keys (`list[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.

        Return:
            tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss,
            logits and labels (each being optional).
        """
        has_labels = False if len(self.label_names) == 0 else all(inputs.get(k) is not None for k in self.label_names)
        # For CLIP-like models capable of returning loss values.
        # If `return_loss` is not specified or being `None` in `inputs`, we check if the default value of `return_loss`
        # is `True` in `model.forward`.
        return_loss = inputs.get("return_loss")
        if return_loss is None:
            return_loss = self.can_return_loss
        loss_without_labels = len(self.label_names) == 0 and return_loss

        inputs = self._prepare_inputs(inputs)
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", ["past_key_values"])
            else:
                ignore_keys = []

        # labels may be popped when computing the loss (label smoothing for instance) so we grab them first.
        if has_labels or loss_without_labels:
            labels = nested_detach(tuple(inputs.get(name) for name in self.label_names))
            if len(labels) == 1:
                labels = labels[0]
        else:
            labels = None

        with torch.no_grad():
            if is_sagemaker_mp_enabled():
                raw_outputs = smp_forward_only(model, inputs)
                if has_labels or loss_without_labels:
                    if isinstance(raw_outputs, dict):
                        loss_mb = raw_outputs["loss"]
                        logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys + ["loss"])
                    else:
                        loss_mb = raw_outputs[0]
                        logits_mb = raw_outputs[1:]

                    loss = loss_mb.reduce_mean().detach().cpu()
                    logits = smp_nested_concat(logits_mb)
                else:
                    loss = None
                    if isinstance(raw_outputs, dict):
                        logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys)
                    else:
                        logits_mb = raw_outputs
                    logits = smp_nested_concat(logits_mb)
            else:
                if has_labels or loss_without_labels:
                    with self.compute_loss_context_manager():
                        num_items_in_batch = self._get_num_items_in_batch([inputs], self.args.device)
                        loss, outputs = self.compute_loss(
                            model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch
                        )
                    loss = loss.detach().mean()

                    if isinstance(outputs, dict):
                        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
                    else:
                        logits = outputs[1:]
                else:
                    loss = None
                    with self.compute_loss_context_manager():
                        outputs = model(**inputs)
                    if isinstance(outputs, dict):
                        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys)
                    else:
                        logits = outputs

        if prediction_loss_only:
            return (loss, None, None)

        logits = nested_detach(logits)
        if len(logits) == 1:
            logits = logits[0]

        return (loss, logits, labels)

    def _evaluate(
        self,
        trial: "optuna.Trial | dict[str, Any] | None",
        ignore_keys_for_eval: list[str] | None,
        skip_scheduler: bool = False,
    ) -> dict[str, float]:
        """Run evaluation, report to HP search, and step ReduceLROnPlateau/GreedyLR if needed."""
        metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
        self._report_to_hp_search(trial, self.state.global_step, metrics)

        # Run delayed LR scheduler now that metrics are populated
        if (
            isinstance(self.lr_scheduler, (torch.optim.lr_scheduler.ReduceLROnPlateau, GreedyLR))
            and not skip_scheduler
        ):
            metric_to_check = self.args.metric_for_best_model
            if not metric_to_check.startswith("eval_"):
                metric_to_check = f"eval_{metric_to_check}"
            try:
                self.lr_scheduler.step(metrics[metric_to_check])
            except KeyError as exc:
                raise KeyError(
                    f"The `metric_for_best_model` training argument is set to '{metric_to_check}', "
                    f"which is not found in the evaluation metrics. "
                    f"The available evaluation metrics are: {list(metrics.keys())}. "
                    f"Please ensure that the `compute_metrics` function returns a dictionary that includes '{metric_to_check}' or "
                    f"consider changing the `metric_for_best_model` via the TrainingArguments."
                ) from exc
        return metrics

    # ---- Checkpoint Saving ----

    def _get_output_dir(self, trial: "optuna.Trial | dict[str, Any] | None") -> str:
        """Return the output directory, accounting for hyperparameter search trials."""
        if self.hp_search_backend is not None and trial is not None:
            if self.hp_search_backend == HPSearchBackend.OPTUNA:
                run_id = trial.number
            elif self.hp_search_backend == HPSearchBackend.RAY:
                import ray.tune

                run_id = ray.tune.get_context().get_trial_id()
            elif self.hp_search_backend == HPSearchBackend.WANDB:
                import wandb

                run_id = wandb.run.id
            run_name = self.hp_name(trial) if self.hp_name is not None else f"run-{run_id}"
            run_dir = os.path.join(self.args.output_dir, run_name)
        else:
            run_dir = self.args.output_dir
        return run_dir

    def _save_checkpoint(self, model: nn.Module, trial: "optuna.Trial | dict[str, Any] | None") -> None:
        """Save model checkpoint, optimizer, scheduler, scaler, RNG states, and trainer state."""
        # In all cases, including ddp/dp/deepspeed, self.model is always a reference to the model we
        # want to save except FullyShardedDDP.
        # assert unwrap_model(model) is self.model, "internal model should be a reference to self.model"

        # Save model checkpoint
        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"

        if self.hp_search_backend is None and trial is None:
            self.store_flos()

        run_dir = self._get_output_dir(trial=trial)
        output_dir = os.path.join(run_dir, checkpoint_folder)
        self.save_model(output_dir, _internal_call=True)

        if (
            self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH, SaveStrategy.BEST]
            and self.state.best_global_step
        ):
            # Wait for everyone to get here so we are sure the model has been saved by process 0
            # before we check if the best_checkpoint_dir exists
            if is_torch_xla_available():
                xm.rendezvous("load_best_model_at_end")
            elif self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                dist.barrier()
            elif is_sagemaker_mp_enabled():
                smp.barrier()

            best_checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.best_global_step}"
            best_checkpoint_dir = os.path.join(run_dir, best_checkpoint_folder)

            if os.path.exists(best_checkpoint_dir):
                self.state.best_model_checkpoint = best_checkpoint_dir

        if not self.args.save_only_model:
            # Save optimizer and scheduler
            self._save_optimizer_and_scheduler(output_dir)
            self._save_scaler(output_dir)
            # Save RNG state
            self._save_rng_state(output_dir)

        # Save the Trainer state
        if self.args.should_save:
            # Update `ExportableState` callbacks and `TrainerControl` state to where we are currently
            for cb in [
                cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
            ]:
                cb_name = cb.__class__.__name__
                cb_state = cb.state()
                if isinstance(self.state.stateful_callbacks[cb_name], list):
                    self.state.stateful_callbacks[cb_name].append(cb_state)
                else:
                    self.state.stateful_callbacks[cb_name] = cb_state
            self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))

        if self.args.push_to_hub:
            self._push_from_checkpoint(output_dir)

        # Maybe delete some older checkpoints.
        if self.args.should_save:
            # we use mtime as default, filesystems without mtime support will be detected in `sort_checkpoints`
            rotate_checkpoints(
                output_dir=run_dir,
                save_total_limit=self.args.save_total_limit,
                best_model_checkpoint=self.state.best_model_checkpoint,
                use_mtime=True,
            )

    def _determine_best_metric(self, metrics: dict[str, float], trial: "optuna.Trial | dict[str, Any] | None") -> bool:
        """
        Determine if the model should be saved based on the evaluation metrics.

        Returns:
            bool: True if a new best metric was found, else False
        """
        is_new_best_metric = False

        if self.args.metric_for_best_model is not None:
            metric_to_check = self.args.metric_for_best_model

            if not metric_to_check.startswith("eval_"):
                metric_to_check = f"eval_{metric_to_check}"

            try:
                metric_value = metrics[metric_to_check]
            except KeyError as exc:
                raise KeyError(
                    f"The `metric_for_best_model` training argument is set to '{metric_to_check}', which is not found in the evaluation metrics. "
                    f"The available evaluation metrics are: {list(metrics.keys())}. Consider changing the `metric_for_best_model` via the TrainingArguments."
                ) from exc

            operator = np.greater if self.args.greater_is_better else np.less

            if self.state.best_metric is None:
                self.state.best_metric = float("-inf") if self.args.greater_is_better else float("inf")

            if operator(metric_value, self.state.best_metric):
                self.state.best_metric = metric_value

                if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH, SaveStrategy.BEST]:
                    self.state.best_global_step = self.state.global_step

                is_new_best_metric = True

        return is_new_best_metric

    def _save_rng_state(self, output_dir: str) -> None:
        """Save random number generator states for reproducible resumption."""
        # Save RNG state in non-distributed training
        rng_states = {
            "python": random.getstate(),
            "numpy": np.random.get_state(),
            "cpu": torch.random.get_rng_state(),
        }
        if torch.cuda.is_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                # In non distributed, we save the global CUDA RNG state (will take care of DataParallel)
                rng_states["cuda"] = torch.cuda.random.get_rng_state_all()
            else:
                rng_states["cuda"] = torch.cuda.random.get_rng_state()

        if is_torch_xla_available():
            rng_states["xla"] = xm.get_rng_state()

        if is_torch_npu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                rng_states["npu"] = torch.npu.random.get_rng_state_all()
            else:
                rng_states["npu"] = torch.npu.random.get_rng_state()

        if is_torch_hpu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                rng_states["hpu"] = torch.hpu.random.get_rng_state_all()
            else:
                rng_states["hpu"] = torch.hpu.random.get_rng_state()

        if is_torch_mlu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                rng_states["mlu"] = torch.mlu.random.get_rng_state_all()
            else:
                rng_states["mlu"] = torch.mlu.random.get_rng_state()

        if is_torch_musa_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                rng_states["musa"] = torch.musa.get_rng_state_all()
            else:
                rng_states["musa"] = torch.musa.get_rng_state()

        # A process can arrive here before the process 0 has a chance to save the model, in which case output_dir may
        # not yet exist.
        os.makedirs(output_dir, exist_ok=True)

        if self.args.world_size <= 1:
            torch.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
        else:
            torch.save(rng_states, os.path.join(output_dir, f"rng_state_{self.args.process_index}.pth"))

    def _save_optimizer_and_scheduler(self, output_dir: str) -> None:
        """Save optimizer and learning rate scheduler states to `output_dir`."""
        if is_torch_xla_available():
            xm.rendezvous("saving_optimizer_states")
            if self.is_fsdp_xla_v1_enabled:
                optm = {
                    "optimizer": self.optimizer.state_dict(),
                    "shard_metadata": self.model.get_shard_metadata(),
                }
                xm.save(
                    optm,
                    os.path.join(
                        output_dir, f"rank{self.args.process_index}-of-{self.args.world_size}-{OPTIMIZER_NAME}"
                    ),
                    master_only=False,
                )
            else:
                xm.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
            with warnings.catch_warnings(record=True) as caught_warnings:
                xm.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
                reissue_pt_warnings(caught_warnings)
        elif is_sagemaker_mp_enabled():
            opt_state_dict = self.optimizer.local_state_dict(gather_if_shard=False)
            smp.barrier()
            if smp.rdp_rank() == 0 or smp.state.cfg.shard_optimizer_state:
                smp.save(
                    opt_state_dict,
                    os.path.join(output_dir, OPTIMIZER_NAME),
                    partial=True,
                    v3=smp.state.cfg.shard_optimizer_state,
                )
        elif self.is_deepspeed_enabled:
            # under zero3 model file itself doesn't get saved since it's bogus! Unless deepspeed
            # config `stage3_gather_16bit_weights_on_model_save` is True
            accept_exclude_frozen_parameters = "exclude_frozen_parameters" in set(
                inspect.signature(self.model_wrapped.save_checkpoint).parameters.keys()
            )
            if accept_exclude_frozen_parameters and _is_peft_model(self.model):
                self.model_wrapped.save_checkpoint(output_dir, exclude_frozen_parameters=True)
            else:
                self.model_wrapped.save_checkpoint(output_dir)
        elif self.is_fsdp_enabled:
            # save fsdp specific ckpt for resuming from ckpt
            save_fsdp_model(
                self.accelerator.state.fsdp_plugin, self.accelerator, self.model, output_dir, **get_fsdp_ckpt_kwargs()
            )
            save_fsdp_optimizer(
                self.accelerator.state.fsdp_plugin, self.accelerator, self.optimizer, self.model, output_dir
            )
        elif self.args.should_save:
            # deepspeed.save_checkpoint above saves model/optim/sched
            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))

        # Save SCHEDULER & SCALER
        is_deepspeed_custom_scheduler = self.is_deepspeed_enabled and not isinstance(
            self.lr_scheduler, DeepSpeedSchedulerWrapper
        )
        if (
            self.args.should_save
            and (not self.is_deepspeed_enabled or is_deepspeed_custom_scheduler)
            and not is_torch_xla_available()
        ):
            with warnings.catch_warnings(record=True) as caught_warnings:
                torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
            reissue_pt_warnings(caught_warnings)

    def _save_scaler(self, output_dir: str) -> None:
        """Save the gradient scaler state if one exists."""
        # See if there is a scaler attribute
        try:
            scaler = self.accelerator.scaler
        except AttributeError:
            return
        if scaler is None:
            return
        if is_torch_xla_available():
            xm.rendezvous("saving_scaler_state")
            with warnings.catch_warnings(record=True) as caught_warnings:
                xm.save(self.accelerator.scaler.state_dict(), os.path.join(output_dir, SCALER_NAME))
                reissue_pt_warnings(caught_warnings)

        # Save SCALER
        if self.args.should_save and not is_torch_xla_available():
            with warnings.catch_warnings(record=True) as caught_warnings:
                torch.save(self.accelerator.scaler.state_dict(), os.path.join(output_dir, SCALER_NAME))
            reissue_pt_warnings(caught_warnings)

    # ---- Checkpoint Resuming ----

    def _load_from_checkpoint(self, resume_from_checkpoint: str, model: nn.Module | None = None) -> None:
        """Load model weights from a checkpoint directory."""
        if model is None:
            model = self.model

        config_file = os.path.join(resume_from_checkpoint, CONFIG_NAME)
        adapter_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_WEIGHTS_NAME)
        adapter_safe_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME)
        weights_file = os.path.join(resume_from_checkpoint, WEIGHTS_NAME)
        weights_index_file = os.path.join(resume_from_checkpoint, WEIGHTS_INDEX_NAME)
        safe_weights_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_NAME)
        safe_weights_index_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_INDEX_NAME)
        is_fsdp_ckpt = os.path.isdir(resume_from_checkpoint) and (
            # this checks the FSDP state dict when `SHARDED_STATE_DICT` is used
            any(
                FSDP_MODEL_NAME in folder_name
                for folder_name in os.listdir(resume_from_checkpoint)
                if os.path.isdir(os.path.join(resume_from_checkpoint, folder_name))
            )
            # this checks the FSDP state dict when `FULL_STATE_DICT` is used
            or os.path.isfile(os.path.join(resume_from_checkpoint, f"{FSDP_MODEL_NAME}.bin"))
        )
        # if multiple adapters exist, they get saved in sub directories
        adapter_subdirs = (
            [
                folder_name
                for folder_name in os.listdir(resume_from_checkpoint)
                if os.path.isdir(os.path.join(resume_from_checkpoint, folder_name))
                and (
                    os.path.isfile(os.path.join(resume_from_checkpoint, folder_name, ADAPTER_WEIGHTS_NAME))
                    or os.path.isfile(os.path.join(resume_from_checkpoint, folder_name, ADAPTER_SAFE_WEIGHTS_NAME))
                )
            ]
            if os.path.isdir(resume_from_checkpoint)
            else []
        )

        if is_fsdp_ckpt and not self.is_fsdp_enabled:
            raise ValueError(f"Checkpoint found at {resume_from_checkpoint} is only supported when using PyTorch FSDP")

        if not (
            any(
                os.path.isfile(f)
                for f in [
                    weights_file,
                    safe_weights_file,
                    weights_index_file,
                    safe_weights_index_file,
                    adapter_weights_file,
                    adapter_safe_weights_file,
                ]
            )
            or is_fsdp_ckpt
            or adapter_subdirs
        ):
            raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")

        logger.info(f"Loading model from {resume_from_checkpoint}.")

        if os.path.isfile(config_file):
            config = PreTrainedConfig.from_json_file(config_file)
            checkpoint_version = config.transformers_version
            if checkpoint_version is not None and checkpoint_version != __version__:
                logger.warning(
                    f"You are resuming training from a checkpoint trained with {checkpoint_version} of "
                    f"Transformers but your current version is {__version__}. This is not recommended and could "
                    "yield to errors or unwanted behaviors."
                )

        if os.path.isfile(weights_file) or os.path.isfile(safe_weights_file) or is_fsdp_ckpt:
            # If the model is on the GPU, it still works!
            if is_sagemaker_mp_enabled():
                smp.resume_from_checkpoint(
                    path=resume_from_checkpoint, tag=WEIGHTS_NAME, partial=False, load_optimizer=False
                )
            elif self.is_fsdp_enabled:
                load_fsdp_model(
                    self.accelerator.state.fsdp_plugin,
                    self.accelerator,
                    model,
                    resume_from_checkpoint,
                    **get_fsdp_ckpt_kwargs(),
                )
            else:
                # We load the model state dict on the CPU to avoid an OOM error.
                if os.path.isfile(safe_weights_file):
                    state_dict = safetensors.torch.load_file(safe_weights_file, device="cpu")
                else:
                    check_torch_load_is_safe()
                    state_dict = torch.load(weights_file, map_location="cpu", weights_only=True)

                # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963
                # which takes *args instead of **kwargs
                load_result = model.load_state_dict(state_dict, False)
                # release memory
                del state_dict
                self._issue_warnings_after_load(load_result)

        # Load adapters following PR # 24096
        elif _is_peft_model(model):
            # If training a model using PEFT, assume that adapter have been saved properly.
            if hasattr(model, "active_adapters") and hasattr(model, "load_adapter"):
                if os.path.exists(resume_from_checkpoint):
                    active_adapters = model.active_adapters
                    if len(active_adapters) > 1:
                        logger.warning("Multiple active adapters detected will only consider the first adapter")
                    active_adapter = active_adapters[0]

                    if adapter_subdirs:
                        for subdir_name in adapter_subdirs:
                            peft_id = os.path.join(resume_from_checkpoint, subdir_name)
                            model.load_adapter(peft_id, subdir_name, is_trainable=(subdir_name == active_adapter))
                        model.set_adapter(active_adapter)
                    else:
                        model.load_adapter(resume_from_checkpoint, active_adapter, is_trainable=True)
                else:
                    logger.warning(
                        "The intermediate checkpoints of PEFT may not be saved correctly, "
                        f"consider using a custom callback to save {ADAPTER_WEIGHTS_NAME} in corresponding saving folders. "
                        "Check some examples here: https://github.com/huggingface/peft/issues/96"
                    )
            else:
                logger.warning(f"Could not load adapter model, make sure to have PEFT >= {MIN_PEFT_VERSION} installed")
        else:
            # We load the sharded checkpoint
            load_result = load_sharded_checkpoint(model, resume_from_checkpoint, strict=is_sagemaker_mp_enabled())
            if not is_sagemaker_mp_enabled():
                self._issue_warnings_after_load(load_result)

    def _load_best_model(self) -> None:
        """Load the best model found during training based on the tracked metric."""
        logger.info(f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric}).")
        best_model_path = os.path.join(self.state.best_model_checkpoint, WEIGHTS_NAME)
        best_safe_model_path = os.path.join(self.state.best_model_checkpoint, SAFE_WEIGHTS_NAME)
        best_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_WEIGHTS_NAME)
        best_safe_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME)

        model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
        if self.is_deepspeed_enabled:
            deepspeed_load_checkpoint(
                self.model_wrapped,
                self.state.best_model_checkpoint,
                load_module_strict=not _is_peft_model(self.model),
            )
        elif self.is_fsdp_enabled:
            load_result = load_fsdp_model(
                self.accelerator.state.fsdp_plugin,
                self.accelerator,
                model,
                self.state.best_model_checkpoint,
                **get_fsdp_ckpt_kwargs(),
            )
        elif (
            os.path.exists(best_model_path)
            or os.path.exists(best_safe_model_path)
            or os.path.exists(best_adapter_model_path)
            or os.path.exists(best_safe_adapter_model_path)
        ):
            has_been_loaded = True
            if is_sagemaker_mp_enabled():
                smp.resume_from_checkpoint(
                    path=self.state.best_model_checkpoint,
                    tag=WEIGHTS_NAME,
                    partial=False,
                    load_optimizer=False,
                )
            else:
                if _is_peft_model(model):
                    # If training a model using PEFT, assume that adapter have been saved properly.
                    if hasattr(model, "active_adapters") and hasattr(model, "load_adapter"):
                        active_adapter = model.active_adapters[0]
                        if len(model.active_adapters) > 1:
                            logger.warning("Detected multiple active adapters, will only consider the first one")

                        if os.path.exists(best_adapter_model_path) or os.path.exists(best_safe_adapter_model_path):
                            try:
                                model.load_adapter(self.state.best_model_checkpoint, active_adapter)
                            except RuntimeError as exc:
                                if model.peft_config[active_adapter].is_prompt_learning:
                                    # for context: https://github.com/huggingface/peft/issues/2256
                                    msg = (
                                        "When using prompt learning PEFT methods such as "
                                        f"{model.peft_config[active_adapter].peft_type.value}, setting "
                                        "load_best_model_at_end=True can lead to errors, it is recommended "
                                        "to set this to False and to load the model manually from the checkpoint "
                                        "directory using PeftModel.from_pretrained(base_model, <path>) after training "
                                        "has finished."
                                    )
                                    raise RuntimeError(msg) from exc
                                else:
                                    raise
                            # Load_adapter has no return value present, modify it when appropriate.
                            from torch.nn.modules.module import _IncompatibleKeys

                            load_result = _IncompatibleKeys([], [])
                        else:
                            logger.warning(
                                "The intermediate checkpoints of PEFT may not be saved correctly, "
                                f"consider using a custom callback to save {ADAPTER_WEIGHTS_NAME} in corresponding saving folders. "
                                "Check some examples here: https://github.com/huggingface/peft/issues/96"
                            )
                            has_been_loaded = False
                    else:
                        logger.warning(
                            f"Could not load adapter model, make sure to have PEFT >= {MIN_PEFT_VERSION} installed"
                        )
                        has_been_loaded = False
                else:
                    # We load the model state dict on the CPU to avoid an OOM error.
                    if os.path.isfile(best_safe_model_path):
                        state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu")
                    else:
                        check_torch_load_is_safe()
                        state_dict = torch.load(best_model_path, map_location="cpu", weights_only=True)

                    # If the model is on the GPU, it still works!
                    # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963
                    # which takes *args instead of **kwargs
                    load_result = model.load_state_dict(state_dict, False)
                if not is_sagemaker_mp_enabled() and has_been_loaded:
                    self._issue_warnings_after_load(load_result)
        elif os.path.exists(os.path.join(self.state.best_model_checkpoint, SAFE_WEIGHTS_INDEX_NAME)) or os.path.exists(
            os.path.join(self.state.best_model_checkpoint, WEIGHTS_INDEX_NAME)
        ):
            load_result = load_sharded_checkpoint(
                model, self.state.best_model_checkpoint, strict=is_sagemaker_mp_enabled()
            )
            if not is_sagemaker_mp_enabled():
                self._issue_warnings_after_load(load_result)
        else:
            logger.warning(
                f"Could not locate the best model at {best_model_path}, if you are running a distributed training "
                "on multiple nodes, you should activate `--save_on_each_node`."
            )

    def _load_rng_state(self, checkpoint: str | None) -> None:
        """Restore random number generator states from a checkpoint."""
        # Load RNG states from `checkpoint`
        if checkpoint is None:
            return

        if self.args.world_size > 1:
            process_index = self.args.process_index
            rng_file = os.path.join(checkpoint, f"rng_state_{process_index}.pth")
            if not os.path.isfile(rng_file):
                logger.info(
                    f"Didn't find an RNG file for process {process_index}, if you are resuming a training that "
                    "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
                )
                return
        else:
            rng_file = os.path.join(checkpoint, "rng_state.pth")
            if not os.path.isfile(rng_file):
                logger.info(
                    "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
                    "fashion, reproducibility is not guaranteed."
                )
                return

        with safe_globals():
            check_torch_load_is_safe()
            checkpoint_rng_state = torch.load(rng_file, weights_only=True)
        random.setstate(checkpoint_rng_state["python"])
        np.random.set_state(checkpoint_rng_state["numpy"])
        torch.random.set_rng_state(checkpoint_rng_state["cpu"])
        if is_torch_xla_available():
            xm.set_rng_state(checkpoint_rng_state["xla"])

        is_distributed = self.args.parallel_mode == ParallelMode.DISTRIBUTED
        if torch.cuda.is_available():
            set_rng_state_for_device("CUDA", torch.cuda, checkpoint_rng_state, is_distributed)
        if is_torch_npu_available():
            set_rng_state_for_device("NPU", torch.npu, checkpoint_rng_state, is_distributed)
        if is_torch_hpu_available():
            set_rng_state_for_device("HPU", torch.hpu, checkpoint_rng_state, is_distributed)
        if is_torch_mlu_available():
            set_rng_state_for_device("MLU", torch.mlu, checkpoint_rng_state, is_distributed)
        if is_torch_musa_available():
            set_rng_state_for_device("MUSA", torch.musa, checkpoint_rng_state, is_distributed)

    def _load_optimizer_and_scheduler(self, checkpoint: str | None) -> None:
        """If optimizer and scheduler states exist, load them."""
        if checkpoint is None:
            return

        if self.is_deepspeed_enabled:
            # deepspeed loads optimizer/lr_scheduler together with the model in deepspeed_init
            if not isinstance(self.lr_scheduler, DeepSpeedSchedulerWrapper):
                with warnings.catch_warnings(record=True) as caught_warnings:
                    check_torch_load_is_safe()
                    self.lr_scheduler.load_state_dict(
                        torch.load(os.path.join(checkpoint, SCHEDULER_NAME), weights_only=True)
                    )
                reissue_pt_warnings(caught_warnings)
            return

        checkpoint_file_exists = (
            glob.glob(os.path.join(checkpoint, OPTIMIZER_NAME) + "_*")
            if is_sagemaker_mp_enabled()
            else (
                os.path.isfile(os.path.join(checkpoint, OPTIMIZER_NAME))
                or os.path.isfile(os.path.join(checkpoint, OPTIMIZER_NAME_BIN))
                or (
                    os.path.isdir(checkpoint)
                    and any(
                        OPTIMIZER_NAME_BIN.split(".")[0] in folder_name
                        for folder_name in os.listdir(checkpoint)
                        if os.path.isdir(os.path.join(checkpoint, folder_name))
                    )
                )
            )
        )
        checkpoint_file_exists = (
            glob.glob(os.path.join(checkpoint, f"rank*-of-{self.args.world_size}-{OPTIMIZER_NAME}"))
            if self.is_fsdp_xla_v1_enabled
            else checkpoint_file_exists
        )
        if checkpoint_file_exists and os.path.isfile(os.path.join(checkpoint, SCHEDULER_NAME)):
            # Load in optimizer and scheduler states
            if is_torch_xla_available():
                # On TPU we have to take some extra precautions to properly load the states on the right device.
                if self.is_fsdp_xla_v1_enabled:
                    check_torch_load_is_safe()
                    optimizer_state = torch.load(
                        os.path.join(
                            checkpoint, f"rank{self.args.process_index}-of-{self.args.world_size}-{OPTIMIZER_NAME}"
                        ),
                        map_location="cpu",
                        weights_only=True,
                    )
                    # We only need `optimizer` when resuming from checkpoint
                    optimizer_state = optimizer_state["optimizer"]
                else:
                    check_torch_load_is_safe()
                    optimizer_state = torch.load(
                        os.path.join(checkpoint, OPTIMIZER_NAME), map_location="cpu", weights_only=True
                    )
                with warnings.catch_warnings(record=True) as caught_warnings:
                    check_torch_load_is_safe()
                    lr_scheduler_state = torch.load(
                        os.path.join(checkpoint, SCHEDULER_NAME), map_location="cpu", weights_only=True
                    )
                reissue_pt_warnings(caught_warnings)

                xm.send_cpu_data_to_device(optimizer_state, self.args.device)
                xm.send_cpu_data_to_device(lr_scheduler_state, self.args.device)

                self.optimizer.load_state_dict(optimizer_state)
                self.lr_scheduler.load_state_dict(lr_scheduler_state)
            else:
                if is_sagemaker_mp_enabled():

                    def opt_load_hook(mod, opt):
                        opt.load_state_dict(smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True))

                    self.model_wrapped.register_post_step_hook(opt_load_hook)
                else:
                    # We use the CPU when training on one GPU to avoid OOM for GPU RAM when training big models.
                    # In distributed training however, we load directly on each GPU and risk the GPU OOM as it's more
                    # likely to get OOM on CPU (since we load num_gpu times the optimizer state
                    map_location = self.args.device if self.args.world_size > 1 else "cpu"
                    if self.is_fsdp_enabled:
                        load_fsdp_optimizer(
                            self.accelerator.state.fsdp_plugin,
                            self.accelerator,
                            self.optimizer,
                            self.model,
                            checkpoint,
                            **get_fsdp_ckpt_kwargs(),
                        )
                    else:
                        check_torch_load_is_safe()
                        self.optimizer.load_state_dict(
                            torch.load(
                                os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location, weights_only=True
                            )
                        )
                with warnings.catch_warnings(record=True) as caught_warnings:
                    check_torch_load_is_safe()
                    self.lr_scheduler.load_state_dict(
                        torch.load(os.path.join(checkpoint, SCHEDULER_NAME), weights_only=True)
                    )
                reissue_pt_warnings(caught_warnings)

    def _load_scaler(self, checkpoint: str | None) -> None:
        """If scaler state exists, load it."""
        if checkpoint is None:
            return

        checkpoint_file_exists = os.path.isfile(os.path.join(checkpoint, SCALER_NAME))

        if checkpoint_file_exists:
            # On TPU we have to take some extra precautions to properly load the states on the right device.
            # Load in scaler states
            if is_torch_xla_available():
                with warnings.catch_warnings(record=True) as caught_warnings:
                    check_torch_load_is_safe()
                    scaler_state = torch.load(
                        os.path.join(checkpoint, SCALER_NAME), map_location="cpu", weights_only=True
                    )
                reissue_pt_warnings(caught_warnings)
                xm.send_cpu_data_to_device(scaler_state, self.args.device)
                self.accelerator.scaler.load_state_dict(scaler_state)
            else:
                with warnings.catch_warnings(record=True) as caught_warnings:
                    check_torch_load_is_safe()
                    self.accelerator.scaler.load_state_dict(
                        torch.load(os.path.join(checkpoint, SCALER_NAME), weights_only=True)
                    )
                reissue_pt_warnings(caught_warnings)

    def _load_callback_state(self) -> None:
        """If callback states exist and were passed in, restore their states if enabled"""
        if not self.args.restore_callback_states_from_checkpoint:
            return
        # Callback states are stored in stateful_callbacks
        not_found = []
        new_callbacks = []
        original_callbacks = self.callback_handler.callbacks + [self.control]
        for stored_callback, data in self.state.stateful_callbacks.items():
            if not isinstance(data, list):
                data = [data]
            if any(callback.__class__.__name__ == stored_callback for callback in original_callbacks):
                # We can load/restore from multiple callbacks of the same type.
                duplicates = [
                    callback for callback in original_callbacks if callback.__class__.__name__ == stored_callback
                ]
                for callback, callback_data in zip(duplicates, data):
                    args = callback_data.get("args", {})
                    attributes = callback_data.get("attributes", {})
                    new_callback = type(callback)(**args)
                    for attribute, value in attributes.items():
                        setattr(new_callback, attribute, value)
                    if isinstance(callback, TrainerControl):
                        # Specifically for restoring the `control` state
                        self.control = new_callback
                    else:
                        new_callbacks.append(new_callback)
                    # We remove the existing callback and add it to the list of new callbacks
                    self.callback_handler.remove_callback(type(new_callback))
                logger.info("Continuing training from checkpoint, restoring any callbacks that were passed in")
            else:
                not_found.append(stored_callback)
        if len(not_found) > 0:
            logger.warning(
                f"Checkpoint included callbacks not included in current configuration. Ignoring. ({', '.join(not_found)})"
            )
        for callback in new_callbacks:
            self.callback_handler.add_callback(callback)

    def _issue_warnings_after_load(self, load_result: Any) -> None:
        """Log warnings for missing or unexpected keys after loading a checkpoint."""
        if len(load_result.missing_keys) != 0:
            if self.model._keys_to_ignore_on_save is not None and set(load_result.missing_keys) == set(
                self.model._keys_to_ignore_on_save
            ):
                self.model.tie_weights()
            else:
                logger.warning(f"There were missing keys in the checkpoint model loaded: {load_result.missing_keys}.")
        if len(load_result.unexpected_keys) != 0:
            logger.warning(
                f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}."
            )

    # ---- Saving & Serialization ----

    def save_model(self, output_dir: str | None = None, _internal_call: bool = False) -> None:
        """
        Will save the model, so you can reload it using `from_pretrained()`.

        Will only save from the main process.
        """

        if output_dir is None:
            output_dir = self.args.output_dir

        if is_torch_xla_available():
            save_tpu_checkpoint(
                self.model, self.args, self.accelerator, self.processing_class, self.is_fsdp_xla_v1_enabled, output_dir
            )
        elif is_sagemaker_mp_enabled():
            # Calling the state_dict needs to be done on the wrapped model and on all processes.
            os.makedirs(output_dir, exist_ok=True)
            state_dict = self.model_wrapped.state_dict()
            if self.args.should_save:
                self._save(output_dir, state_dict=state_dict)
            Path(os.path.join(output_dir, "user_content.pt")).touch()
        elif self.is_fsdp_enabled:
            if "FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type):
                state_dict = self.accelerator.get_state_dict(self.model)
                if self.args.should_save:
                    self._save(output_dir, state_dict=state_dict)
        elif self.is_deepspeed_enabled:
            try:
                accept_exclude_frozen_parameters = "exclude_frozen_parameters" in set(
                    inspect.signature(self.model_wrapped.save_checkpoint).parameters.keys()
                )
                zero3_sharding = self.deepspeed.config.get("zero_optimization", {}).get("stage", None) == 3
                if accept_exclude_frozen_parameters and _is_peft_model(self.model) and zero3_sharding:
                    # When using PEFT with DeepSpeed ZeRO Stage 3,
                    # we do not need to load the frozen parameters
                    state_dict = self.deepspeed._zero3_consolidated_16bit_state_dict(exclude_frozen_parameters=True)
                else:
                    state_dict = self.accelerator.get_state_dict(self.deepspeed)
                if self.args.should_save:
                    self._save(output_dir, state_dict=state_dict)
            except ValueError:
                logger.warning(
                    " stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead, use"
                    " zero_to_fp32.py to recover weights"
                )
                if self.args.should_save:
                    self._save(output_dir, state_dict={})
                # remove the dummy state_dict
                remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
                self.model_wrapped.save_checkpoint(output_dir)

        elif self.args.should_save:
            self._save(output_dir)

        # Push to the Hub when `save_model` is called by the user.
        if self.args.push_to_hub and not _internal_call:
            self.push_to_hub(commit_message="Model save", revision=self.args.hub_revision)

    def _save(self, output_dir: str | None = None, state_dict: dict | None = None) -> None:
        """Save model weights, configuration, and processing class to `output_dir`."""
        # If we are executing this function, we are the process zero, so we don't check for that.
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Saving model checkpoint to {output_dir}")

        supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel)
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        if not isinstance(self.model, supported_classes):
            if state_dict is None:
                state_dict = self.model.state_dict()

            if isinstance(self.accelerator.unwrap_model(self.model, keep_torch_compile=False), supported_classes):
                self.accelerator.unwrap_model(self.model, keep_torch_compile=False).save_pretrained(
                    output_dir, state_dict=state_dict
                )
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                safetensors.torch.save_file(
                    state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"}
                )
        else:
            self.model.save_pretrained(output_dir, state_dict=state_dict)

        if self.processing_class is not None:
            self.processing_class.save_pretrained(output_dir)
        elif (
            self.data_collator is not None
            and hasattr(self.data_collator, "tokenizer")
            and self.data_collator.tokenizer is not None
        ):
            logger.info("Saving Trainer.data_collator.tokenizer by default as Trainer.processing_class is `None`")
            self.data_collator.tokenizer.save_pretrained(output_dir)

        # Good practice: save your training arguments together with the trained model
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

    # ---- Logging & Metrics ----

    def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
        """
        Log `logs` on the various objects watching training.

        Subclass and override this method to inject custom behavior.

        Args:
            logs (`dict[str, float]`):
                The values to log.
            start_time (`Optional[float]`):
                The start of training.
        """
        if self.state.epoch is not None:
            logs["epoch"] = self.state.epoch
        if self.args.include_num_input_tokens_seen != "no":
            logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen
            if start_time is not None:
                current_session_num_tokens = self.state.num_input_tokens_seen - self._initial_num_input_tokens_seen
                logs.update(speed_metrics("train", start_time, num_tokens=current_session_num_tokens))

        output = {**logs, "step": self.state.global_step}
        self.state.log_history.append(output)
        self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)

    def store_flos(self) -> None:
        """Store the number of floating-point operations that went into the model."""
        if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
            self.state.total_flos += (
                distributed_broadcast_scalars([self.current_flos], device=self.args.device).sum().item()
            )
            self.current_flos = 0
        else:
            self.state.total_flos += self.current_flos
            self.current_flos = 0

    def floating_point_ops(self, inputs: dict[str, torch.Tensor | Any]) -> int:
        """
        For models that inherit from [`PreTrainedModel`], uses that method to compute the number of floating point
        operations for every backward + forward pass. If using another model, either implement such a method in the
        model or subclass and override this method.

        Args:
            inputs (`dict[str, torch.Tensor | Any]`):
                The inputs and targets of the model.

        Returns:
            `int`: The number of floating-point operations.
        """
        if (main_input := getattr(self.model, "main_input_name", "input_ids")) in inputs and hasattr(
            self.model, "num_parameters"
        ):
            return 6 * inputs[main_input].numel() * self.model.num_parameters(exclude_embeddings=True)
        return 0

    # ---- Hub Integration ----

    def init_hf_repo(self, token: str | None = None) -> None:
        """
        Initializes a git repo in `self.args.hub_model_id`.
        """
        # Only on process zero
        if not self.is_world_process_zero():
            return

        if self.args.hub_model_id is None:
            repo_name = Path(self.args.output_dir).absolute().name
        else:
            repo_name = self.args.hub_model_id

        token = token if token is not None else self.args.hub_token
        repo_url = create_repo(repo_name, token=token, private=self.args.hub_private_repo, exist_ok=True)
        self.hub_model_id = repo_url.repo_id
        self.push_in_progress = None

    def create_model_card(
        self,
        language: str | None = None,
        license: str | None = None,
        tags: str | list[str] | None = None,
        model_name: str | None = None,
        finetuned_from: str | None = None,
        tasks: str | list[str] | None = None,
        dataset_tags: str | list[str] | None = None,
        dataset: str | list[str] | None = None,
        dataset_args: str | list[str] | None = None,
    ) -> None:
        """
        Creates a draft of a model card using the information available to the `Trainer`.

        Args:
            language (`str`, *optional*):
                The language of the model (if applicable)
            license (`str`, *optional*):
                The license of the model. Will default to the license of the pretrained model used, if the original
                model given to the `Trainer` comes from a repo on the Hub.
            tags (`str` or `list[str]`, *optional*):
                Some tags to be included in the metadata of the model card.
            model_name (`str`, *optional*):
                The name of the model.
            finetuned_from (`str`, *optional*):
                The name of the model used to fine-tune this one (if applicable). Will default to the name of the repo
                of the original model given to the `Trainer` (if it comes from the Hub).
            tasks (`str` or `list[str]`, *optional*):
                One or several task identifiers, to be included in the metadata of the model card.
            dataset_tags (`str` or `list[str]`, *optional*):
                One or several dataset tags, to be included in the metadata of the model card.
            dataset (`str` or `list[str]`, *optional*):
                One or several dataset identifiers, to be included in the metadata of the model card.
            dataset_args (`str` or `list[str]`, *optional*):
               One or several dataset arguments, to be included in the metadata of the model card.
        """
        if not self.is_world_process_zero():
            return

        model_card_filepath = os.path.join(self.args.output_dir, "README.md")
        is_peft_library = False
        if os.path.exists(model_card_filepath):
            library_name = ModelCard.load(model_card_filepath).data.get("library_name")
            is_peft_library = library_name == "peft"

            # Append existing tags in `tags`
            existing_tags = ModelCard.load(model_card_filepath).data.tags
            if tags is not None and existing_tags is not None:
                if isinstance(tags, str):
                    tags = [tags]
                for tag in existing_tags:
                    if tag not in tags:
                        tags.append(tag)

        training_summary = TrainingSummary.from_trainer(
            self,
            language=language,
            license=license,
            tags=tags,
            model_name=model_name,
            finetuned_from=finetuned_from,
            tasks=tasks,
            dataset_tags=dataset_tags,
            dataset=dataset,
            dataset_args=dataset_args,
        )
        model_card = training_summary.to_model_card()
        with open(model_card_filepath, "w") as f:
            f.write(model_card)

        if is_peft_library:
            self.accelerator.unwrap_model(self.model).create_or_update_model_card(self.args.output_dir)

    def push_to_hub(
        self,
        commit_message: str | None = "End of training",
        blocking: bool = True,
        token: str | None = None,
        revision: str | None = None,
        **kwargs,
    ) -> CommitInfo:
        """
        Upload `self.model` and `self.processing_class` to the 🤗 model hub on the repo `self.args.hub_model_id`.

        Parameters:
            commit_message (`str`, *optional*, defaults to `"End of training"`):
                Message to commit while pushing.
            blocking (`bool`, *optional*, defaults to `True`):
                Whether the function should return only when the `git push` has finished.
            token (`str`, *optional*, defaults to `None`):
                Token with write permission to overwrite Trainer's original args.
            revision (`str`, *optional*):
                The git revision to commit from. Defaults to the head of the "main" branch.
            kwargs (`dict[str, Any]`, *optional*):
                Additional keyword arguments passed along to [`~Trainer.create_model_card`].

        Returns:
            The URL of the repository where the model was pushed if `blocking=False`, or a `Future` object tracking the
            progress of the commit if `blocking=True`.
        """
        self.callback_handler.on_push_begin(self.args, self.state, self.control)

        model_name = kwargs.pop("model_name", None)
        if model_name is None and self.args.should_save:
            if self.args.hub_model_id is None:
                model_name = Path(self.args.output_dir).name
            else:
                model_name = self.args.hub_model_id.split("/")[-1]
        token = token if token is not None else self.args.hub_token

        # In case the user calls this method with args.push_to_hub = False
        if self.hub_model_id is None:
            self.init_hf_repo(token=token)

        # Needs to be executed on all processes for TPU training, but will only save on the processed determined by
        # self.args.should_save.
        self.save_model(_internal_call=True)

        # Only push from one node.
        if not self.is_world_process_zero():
            return

        # Add additional tags in the case the model has already some tags and users pass
        # "tags" argument to `push_to_hub` so that trainer automatically handles internal tags
        # from all models since Trainer does not call `model.push_to_hub`.
        if getattr(self.model, "model_tags", None) is not None:
            if "tags" not in kwargs:
                kwargs["tags"] = []

            # If it is a string, convert it to a list
            if isinstance(kwargs["tags"], str):
                kwargs["tags"] = [kwargs["tags"]]

            for model_tag in self.model.model_tags:
                if model_tag not in kwargs["tags"]:
                    kwargs["tags"].append(model_tag)

        self.create_model_card(model_name=model_name, **kwargs)

        if revision is None:
            revision = self.args.hub_revision

        # Wait for the current upload to be finished.
        self._finish_current_push()

        return upload_folder(
            repo_id=self.hub_model_id,
            folder_path=self.args.output_dir,
            commit_message=commit_message,
            token=token,
            run_as_future=not blocking,
            ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"],
            revision=revision,
        )

    def _push_from_checkpoint(self, checkpoint_folder: str) -> None:
        """Push model and checkpoint files to the Hub from a checkpoint folder."""
        if not self.is_world_process_zero() or self.args.hub_strategy == HubStrategy.END:
            return
        # If we haven't finished the last push, we don't do this one unless args.hub_always_push=True.
        if not self.args.hub_always_push and self.push_in_progress is not None and not self.push_in_progress.is_done():
            return

        self.callback_handler.on_push_begin(self.args, self.state, self.control)
        output_dir = self.args.output_dir
        # To avoid a new synchronization of all model weights, we just copy the file from the checkpoint folder
        modeling_files = [CONFIG_NAME, GENERATION_CONFIG_NAME, WEIGHTS_NAME, SAFE_WEIGHTS_NAME]
        #  Add sharded checkpoints if we have an index
        for index_file in [WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME]:
            index_path = os.path.join(checkpoint_folder, index_file)
            if os.path.isfile(index_path):
                modeling_files.append(index_file)
                with open(index_path) as f:
                    index = json.loads(f.read())
                shard_files = list(set(index["weight_map"].values()))
                modeling_files.extend(shard_files)
        if is_peft_available():
            modeling_files.extend([ADAPTER_CONFIG_NAME, ADAPTER_WEIGHTS_NAME, ADAPTER_SAFE_WEIGHTS_NAME])
        for modeling_file in modeling_files:
            if os.path.isfile(os.path.join(checkpoint_folder, modeling_file)):
                shutil.copy(os.path.join(checkpoint_folder, modeling_file), os.path.join(output_dir, modeling_file))
        # Saving the processing class is fast and we don't know how many files it may have spawned, so we resave it to be sure.
        if self.processing_class is not None:
            self.processing_class.save_pretrained(output_dir)
        # Same for the training arguments
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

        if self.args.save_strategy == SaveStrategy.STEPS:
            commit_message = f"Training in progress, step {self.state.global_step}"
        else:
            commit_message = f"Training in progress, epoch {int(self.state.epoch)}"

        model_push_job = upload_folder(
            repo_id=self.hub_model_id,
            folder_path=output_dir,
            commit_message=commit_message,
            token=self.args.hub_token,
            run_as_future=True,
            ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"],
            revision=self.args.hub_revision,
        )

        push_jobs = [model_push_job]

        if self.args.hub_strategy in [HubStrategy.CHECKPOINT, HubStrategy.ALL_CHECKPOINTS]:
            path_in_repo = (
                "last-checkpoint" if self.args.hub_strategy == HubStrategy.CHECKPOINT else Path(checkpoint_folder).name
            )
            checkpoint_push = upload_folder(
                repo_id=self.hub_model_id,
                folder_path=checkpoint_folder,
                path_in_repo=path_in_repo,
                commit_message=commit_message + ", checkpoint",
                token=self.args.hub_token,
                run_as_future=True,
                revision=self.args.hub_revision,
            )
            push_jobs.append(checkpoint_push)

        if self.push_in_progress is None or self.push_in_progress.is_done():
            self.push_in_progress = PushInProgress(push_jobs)
        else:
            self.push_in_progress.jobs.extend(push_jobs)

    def _finish_current_push(self) -> None:
        """Wait for any in-progress push to the Hub to complete."""
        if not hasattr(self, "push_in_progress"):
            return
        if self.push_in_progress is not None and not self.push_in_progress.is_done():
            logger.info("Waiting for the current checkpoint push to be finished, this might take a couple of minutes.")
            self.push_in_progress.wait_until_done()

    # ---- Hyperparameter Search ----

    def hyperparameter_search(
        self,
        hp_space: Callable[["optuna.Trial"], dict[str, float]] | None = None,
        compute_objective: Callable[[dict[str, float]], float] | None = None,
        n_trials: int = 20,
        direction: str | list[str] = "minimize",
        backend: str | HPSearchBackend | None = None,
        hp_name: Callable[["optuna.Trial"], str] | None = None,
        **kwargs,
    ) -> BestRun | list[BestRun]:
        """
        Launch a hyperparameter search using `optuna` or `Ray Tune`. The optimized quantity is determined
        by `compute_objective`, which defaults to a function returning the evaluation loss when no metric is provided,
        the sum of all metrics otherwise.

        <Tip warning={true}>

        To use this method, you need to have provided a `model_init` when initializing your [`Trainer`]: we need to
        reinitialize the model at each new run. This is incompatible with the `optimizers` argument, so you need to
        subclass [`Trainer`] and override the method [`~Trainer.create_optimizer_and_scheduler`] for custom
        optimizer/scheduler.

        </Tip>

        Args:
            hp_space (`Callable[["optuna.Trial"], dict[str, float]]`, *optional*):
                A function that defines the hyperparameter search space. Will default to
                [`~trainer_utils.default_hp_space_optuna`] or [`~trainer_utils.default_hp_space_ray`]
                depending on your backend.
            compute_objective (`Callable[[dict[str, float]], float]`, *optional*):
                A function computing the objective to minimize or maximize from the metrics returned by the `evaluate`
                method. Will default to [`~trainer_utils.default_compute_objective`].
            n_trials (`int`, *optional*, defaults to 100):
                The number of trial runs to test.
            direction (`str` or `list[str]`, *optional*, defaults to `"minimize"`):
                If it's single objective optimization, direction is `str`, can be `"minimize"` or `"maximize"`, you
                should pick `"minimize"` when optimizing the validation loss, `"maximize"` when optimizing one or
                several metrics. If it's multi objectives optimization, direction is `list[str]`, can be List of
                `"minimize"` and `"maximize"`, you should pick `"minimize"` when optimizing the validation loss,
                `"maximize"` when optimizing one or several metrics.
            backend (`str` or [`~training_utils.HPSearchBackend`], *optional*):
                The backend to use for hyperparameter search. Will default to optuna or Ray Tune, depending
                on which one is installed. If all are installed, will default to optuna.
            hp_name (`Callable[["optuna.Trial"], str]]`, *optional*):
                A function that defines the trial/run name. Will default to None.
            kwargs (`dict[str, Any]`, *optional*):
                Additional keyword arguments for each backend:

                - `optuna`: parameters from
                  [optuna.study.create_study](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html)
                  and also the parameters `timeout`, `n_jobs` and `gc_after_trial` from
                  [optuna.study.Study.optimize](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.optimize)
                - `ray`: parameters from [tune.run](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run).
                  If `resources_per_trial` is not set in the `kwargs`, it defaults to 1 CPU core and 1 GPU (if available).
                  If `progress_reporter` is not set in the `kwargs`,
                  [ray.tune.CLIReporter](https://docs.ray.io/en/latest/tune/api/doc/ray.tune.CLIReporter.html) is used.
        Returns:
            [`trainer_utils.BestRun` or `list[trainer_utils.BestRun]`]: All the information about the best run or best
            runs for multi-objective optimization. Experiment summary can be found in `run_summary` attribute for Ray
            backend.
        """
        if backend is None:
            backend = default_hp_search_backend()
        backend = HPSearchBackend(backend)
        backend_obj = ALL_HYPERPARAMETER_SEARCH_BACKENDS[backend]()
        backend_obj.ensure_available()
        self.hp_search_backend = backend
        if self.model_init is None:
            raise RuntimeError(
                "To use hyperparameter search, you need to pass your model through a model_init function."
            )

        self.hp_space = backend_obj.default_hp_space if hp_space is None else hp_space
        self.hp_name = hp_name
        self.compute_objective = default_compute_objective if compute_objective is None else compute_objective

        best_run = backend_obj.run(self, n_trials, direction, **kwargs)

        self.hp_search_backend = None
        return best_run

    def call_model_init(self, trial: "optuna.Trial | dict[str, Any] | None" = None) -> nn.Module:
        """Invoke `model_init` to get a fresh model instance, optionally conditioned on a hyperparameter trial."""
        model_init_argcount = number_of_arguments(self.model_init)
        if model_init_argcount == 0:
            model = self.model_init()
        elif model_init_argcount == 1:
            model = self.model_init(trial)
        else:
            raise RuntimeError("model_init should have 0 or 1 argument.")

        if model is None:
            raise RuntimeError("model_init should not return None.")

        return model

    def _hp_search_setup(self, trial: "optuna.Trial | dict[str, Any] | None") -> None:
        """Set up training arguments and accelerator state for a hyperparameter search trial."""
        self._trial = trial

        if self.hp_search_backend is None or trial is None:
            return
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            params = self.hp_space(trial)
        elif self.hp_search_backend == HPSearchBackend.RAY:
            params = trial
            params.pop("wandb", None)
        elif self.hp_search_backend == HPSearchBackend.WANDB:
            params = trial

        for key, value in params.items():
            if not hasattr(self.args, key):
                logger.warning(
                    f"Trying to set {key} in the hyperparameter search but there is no corresponding field in"
                    " `TrainingArguments`."
                )
                continue
            old_attr = getattr(self.args, key, None)
            # Casting value to the proper type
            if old_attr is not None:
                value = type(old_attr)(value)

            setattr(self.args, key, value)
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            logger.info(f"Trial: {trial.params}")
        if self.hp_search_backend == HPSearchBackend.WANDB:
            logger.info(f"W&B Sweep parameters: {trial}")
        if self.is_deepspeed_enabled:
            if self.args.deepspeed is None:
                raise ValueError("For sweeps with deepspeed, `args.deepspeed` must be set")

            self.accelerator.free_memory()

            # Rebuild the deepspeed config to reflect the updated training parameters
            from accelerate.utils import DeepSpeedPlugin

            from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig

            self.args.hf_deepspeed_config = HfTrainerDeepSpeedConfig(self.args.deepspeed)
            self.args.hf_deepspeed_config.trainer_config_process(self.args)
            self.args.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.args.hf_deepspeed_config)

            # From 1.0 on, we need to fully wipe the DS plugin when doing sweeps.
            # Simply calling `_reset_state` is enough and doesn't need a version pin.
            AcceleratorState()._reset_state()

        # `train_batch_size` might change when using HPO https://github.com/huggingface/transformers/pull/18918
        self._train_batch_size = self.args.train_batch_size
        self.create_accelerator_and_postprocess()

    def _report_to_hp_search(
        self, trial: "optuna.Trial | dict[str, Any] | None", step: int, metrics: dict[str, float]
    ) -> None:
        """Report intermediate metrics to the active hyperparameter search backend."""
        if self.hp_search_backend is None or trial is None:
            return
        metrics = metrics.copy()
        self.objective = self.compute_objective(metrics)
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            import optuna

            if hasattr(trial, "study") and not trial.study._is_multi_objective():
                trial.report(self.objective, step)
                if trial.should_prune():
                    self.callback_handler.on_train_end(self.args, self.state, self.control)
                    raise optuna.TrialPruned()
        elif self.hp_search_backend == HPSearchBackend.RAY:
            import ray.tune

            with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
                checkpoint = None
                if self.control.should_save:
                    self._tune_save_checkpoint(checkpoint_dir=temp_checkpoint_dir)
                    checkpoint = ray.tune.Checkpoint.from_directory(temp_checkpoint_dir)
                metrics["objective"] = self.objective
                ray.tune.report(metrics, checkpoint=checkpoint)

    def _tune_save_checkpoint(self, checkpoint_dir: str) -> None:
        """Save a checkpoint during a Ray Tune hyperparameter search trial."""
        output_dir = os.path.join(checkpoint_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}")
        self.save_model(output_dir, _internal_call=True)
        if self.args.should_save:
            # Update the `TrainerControl` state to where we are currently
            self.state.stateful_callbacks["TrainerControl"] = self.control.state()
            self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
            torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))

    # ---- Callbacks ----

    def add_callback(self, callback: type[TrainerCallback] | TrainerCallback) -> None:
        """
        Add a callback to the current list of [`~transformers.TrainerCallback`].

        Args:
           callback (`type` or [`~transformers.TrainerCallback]`):
               A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
               first case, will instantiate a member of that class.
        """
        self.callback_handler.add_callback(callback)

    def pop_callback(self, callback: type[TrainerCallback] | TrainerCallback) -> TrainerCallback | None:
        """
        Remove a callback from the current list of [`~transformers.TrainerCallback`] and returns it.

        If the callback is not found, returns `None` (and no error is raised).

        Args:
           callback (`type` or [`~transformers.TrainerCallback]`):
               A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
               first case, will pop the first member of that class found in the list of callbacks.

        Returns:
            [`~transformers.TrainerCallback`]: The callback removed, if found.
        """
        return self.callback_handler.pop_callback(callback)

    def remove_callback(self, callback: type[TrainerCallback] | TrainerCallback) -> None:
        """
        Remove a callback from the current list of [`~transformers.TrainerCallback`].

        Args:
           callback (`type` or [`~transformers.TrainerCallback]`):
               A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
               first case, will remove the first member of that class found in the list of callbacks.
        """
        self.callback_handler.remove_callback(callback)

    # ---- Utilities ----

    def is_local_process_zero(self) -> bool:
        """
        Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several
        machines) main process.
        """
        return self.args.local_process_index == 0

    def is_world_process_zero(self) -> bool:
        """
        Whether or not this process is the global main process (when training in a distributed fashion on several
        machines, this is only going to be `True` for one process).
        """
        # Special case for SageMaker ModelParallel since there process_index is dp_process_index, not the global
        # process index.
        if is_sagemaker_mp_enabled():
            return smp.rank() == 0
        return self.args.process_index == 0

    def _move_model_to_device(self, model: nn.Module, device: torch.device) -> None:
        """Move the model to the specified device, re-tying weights on XLA if needed."""
        if getattr(model, "hf_device_map", None) is not None:
            logger.warning(
                "The model is already on multiple devices. Skipping the move to device specified in `args`."
            )
            return
        model = model.to(device)
        # Moving a model to an XLA device disconnects the tied weights, so we have to retie them.
        if self.args.parallel_mode == ParallelMode.TPU and hasattr(model, "tie_weights"):
            model.tie_weights()

Trainer is a simple but feature-complete training and eval loop for PyTorch, optimized for 🤗 Transformers.

Args -----= model ([PreTrainedModel] or torch.nn.Module, optional): The model to train, evaluate or use for predictions. If not provided, a model_init must be passed.

<Tip>

[`Trainer`] is optimized to work with the [`PreTrainedModel`] provided by the library. You can still use
your own models defined as <code>torch.nn.Module</code> as long as they work the same way as the 🤗 Transformers
models.

</Tip>

args ([TrainingArguments], optional): The arguments to tweak for training. Will default to a basic instance of [TrainingArguments] with the output_dir set to a directory named tmp_trainer in the current directory if not provided. data_collator (DataCollator, optional): The function to use to form a batch from a list of elements of train_dataset or eval_dataset. Will default to [default_data_collator] if no processing_class is provided, an instance of [DataCollatorWithPadding] otherwise if the processing_class is a feature extractor or tokenizer. train_dataset (torch.utils.data.Dataset | torch.utils.data.IterableDataset | datasets.Dataset, optional): The dataset to use for training. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed.

Note that if it's a <code>torch.utils.data.IterableDataset</code> with some randomization and you are training in a
distributed fashion, your iterable dataset should either use a internal attribute <code>generator</code> that is a
<code>torch.Generator</code> for the randomization that must be identical on all processes (and the Trainer will
manually set the seed of this <code>generator</code> at each epoch) or have a <code>set\_epoch()</code> method that internally
sets the seed of the RNGs used.

eval_dataset (torch.utils.data.Dataset | dict[str, torch.utils.data.Dataset] | datasets.Dataset, optional): The dataset to use for evaluation. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed. If it is a dictionary, it will evaluate on each dataset prepending the dictionary key to the metric name. processing_class (PreTrainedTokenizerBase or BaseImageProcessor or FeatureExtractionMixin or ProcessorMixin, optional): Processing class used to process the data. If provided, will be used to automatically process the inputs for the model, and it will be saved along the model to make it easier to rerun an interrupted training or reuse the fine-tuned model. model_init (Callable[[], PreTrainedModel], optional): A function that instantiates the model to be used. If provided, each call to [~Trainer.train] will start from a new instance of the model as given by this function.

The function may have zero argument, or a single one containing the optuna/Ray Tune trial object, to
be able to choose different architectures according to hyperparameters (such as layer count, sizes of
inner layers, dropout probabilities etc).

compute_loss_func (Callable, optional): A function that accepts the raw model outputs, labels, and the number of items in the entire accumulated batch (batch_size * gradient_accumulation_steps) and returns the loss. For example, see the default loss function used by [Trainer]. compute_metrics (Callable[[EvalPrediction], Dict], optional): The function that will be used to compute metrics at evaluation. Must take a [EvalPrediction] and return a dictionary string to metric values. Note When passing TrainingArgs with batch_eval_metrics set to True, your compute_metrics function must take a boolean compute_result argument. This will be triggered after the last eval batch to signal that the function needs to calculate and return the global summary statistics rather than accumulating the batch-level statistics callbacks (List of [TrainerCallback], optional): A list of callbacks to customize the training loop. Will add those to the list of default callbacks detailed in here.

If you want to remove one of the default callbacks used, use the [`Trainer.remove_callback`] method.

optimizers (tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR], optional, defaults to (None, None)): A tuple containing the optimizer and the scheduler to use. Will default to an instance of [AdamW] on your model and a scheduler given by [get_linear_schedule_with_warmup] controlled by args. optimizer_cls_and_kwargs (tuple[Type[torch.optim.Optimizer], dict[str, Any]], optional): A tuple containing the optimizer class and keyword arguments to use. Overrides optim and optim_args in args. Incompatible with the optimizers argument.

Unlike <code>optimizers</code>, this argument avoids the need to place model parameters on the correct devices before initializing the Trainer.

preprocess_logits_for_metrics (Callable[[torch.Tensor, torch.Tensor], torch.Tensor], optional): A function that preprocess the logits right before caching them at each evaluation step. Must take two tensors, the logits and the labels, and return the logits once processed as desired. The modifications made by this function will be reflected in the predictions received by compute_metrics.

Note that the labels (second parameter) will be <code>None</code> if the dataset does not have them.

Important attributes:

- **model** -- Always points to the core model. If using a transformers model, it will be a [`PreTrainedModel`]
  subclass.
- **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the
  original model. This is the model that should be used for the forward pass. For example, under <code>DeepSpeed</code>,
  the inner model is wrapped in <code>DeepSpeed</code> and then again in <code>torch.nn.DistributedDataParallel</code>. If the inner
  model hasn't been wrapped, then <code>self.model\_wrapped</code> is the same as <code>self.model</code>.
- **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from
  data parallelism, this means some of the model layers are split on different GPUs).
- **place_model_on_device** -- Whether or not to automatically place the model on the device. Defaults to
  <code>True</code> unless model parallel, DeepSpeed, FSDP, full fp16/bf16 eval, or SageMaker MP is active. Can be
  overridden by subclassing <code><a title="neuralops.attentions.TrainingArguments" href="#neuralops.attentions.TrainingArguments">TrainingArguments</a></code> and overriding the <code>place\_model\_on\_device</code> property.
- **is_in_train** -- Whether or not a model is currently running <code>train</code> (e.g. when <code>evaluate</code> is called while
  in <code>train</code>)

Subclasses

• MethodTrainer
• SVDAwareTrainer

Static methods

def get_optimizer_cls_and_kwargs(args: transformers.training_args.TrainingArguments,
model: transformers.modeling_utils.PreTrainedModel | None = None) ‑> tuple[typing.Any, typing.Any]

Expand source code

@staticmethod
def get_optimizer_cls_and_kwargs(args: TrainingArguments, model: PreTrainedModel | None = None) -> tuple[Any, Any]:
    """
    Returns the optimizer class and optimizer parameters based on the training arguments.

    Args:
        args (`transformers.training_args.TrainingArguments`):
            The training arguments for the training session.
        model (`PreTrainedModel`, *optional*):
            The model being trained. Required for some optimizers (GaLore, Apollo, LOMO).

    Returns:
        A tuple containing the optimizer class and a dictionary of optimizer keyword arguments.
    """
    ctx = OptimizerContext(
        args=args,
        model=model,
        optimizer_kwargs={"lr": args.learning_rate},
        adam_kwargs={
            "betas": (args.adam_beta1, args.adam_beta2),
            "eps": args.adam_epsilon,
        },
        optim_args=_parse_optim_args(args.optim_args),
    )

    handler = _OPTIMIZER_HANDLERS.get(args.optim)
    if handler is None:
        raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {args.optim}")

    return handler(ctx)

Returns the optimizer class and optimizer parameters based on the training arguments.

Args -----= args (transformers.training_args.TrainingArguments): The training arguments for the training session. model (PreTrainedModel, optional): The model being trained. Required for some optimizers (GaLore, Apollo, LOMO).

Returns -----= A tuple containing the optimizer class and a dictionary of optimizer keyword arguments.

Methods

def add_callback(self,
callback: type[transformers.trainer_callback.TrainerCallback] | transformers.trainer_callback.TrainerCallback) ‑> None

Expand source code

def add_callback(self, callback: type[TrainerCallback] | TrainerCallback) -> None:
    """
    Add a callback to the current list of [`~transformers.TrainerCallback`].

    Args:
       callback (`type` or [`~transformers.TrainerCallback]`):
           A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
           first case, will instantiate a member of that class.
    """
    self.callback_handler.add_callback(callback)

Add a callback to the current list of [~transformers.TrainerCallback].

Args -----= callback (type or [~transformers.TrainerCallback]): A [~transformers.TrainerCallback] class or an instance of a [~transformers.TrainerCallback]. In the first case, will instantiate a member of that class.

def autocast_smart_context_manager(self, cache_enabled: bool | None = True) ‑> contextlib.AbstractContextManager

Expand source code

def autocast_smart_context_manager(self, cache_enabled: bool | None = True) -> contextlib.AbstractContextManager:
    """
    A helper wrapper that creates an appropriate context manager for `autocast` while feeding it the desired
    arguments, depending on the situation. We rely on accelerate for autocast, hence we do nothing here.
    """
    return contextlib.nullcontext()

A helper wrapper that creates an appropriate context manager for autocast while feeding it the desired arguments, depending on the situation. We rely on accelerate for autocast, hence we do nothing here.

def call_model_init(self,
trial: optuna.Trial | dict[str, Any] | None = None) ‑> torch.nn.modules.module.Module

Expand source code

def call_model_init(self, trial: "optuna.Trial | dict[str, Any] | None" = None) -> nn.Module:
    """Invoke `model_init` to get a fresh model instance, optionally conditioned on a hyperparameter trial."""
    model_init_argcount = number_of_arguments(self.model_init)
    if model_init_argcount == 0:
        model = self.model_init()
    elif model_init_argcount == 1:
        model = self.model_init(trial)
    else:
        raise RuntimeError("model_init should have 0 or 1 argument.")

    if model is None:
        raise RuntimeError("model_init should not return None.")

    return model

Invoke model_init to get a fresh model instance, optionally conditioned on a hyperparameter trial.

def compute_loss(self,
model: torch.nn.modules.module.Module,
inputs: dict[str, torch.Tensor | Any],
return_outputs: bool = False,
num_items_in_batch: torch.Tensor | int | None = None) ‑> torch.Tensor | tuple[torch.Tensor, typing.Any]

Expand source code

def compute_loss(
    self,
    model: nn.Module,
    inputs: dict[str, torch.Tensor | Any],
    return_outputs: bool = False,
    num_items_in_batch: torch.Tensor | int | None = None,
) -> torch.Tensor | tuple[torch.Tensor, Any]:
    """
    How the loss is computed by Trainer. By default, all models return the loss in the first element.

    Args:
        model (`nn.Module`):
            The model to compute the loss for.
        inputs (`dict[str, torch.Tensor | Any]`):
            The input data for the model.
        return_outputs (`bool`, *optional*, defaults to `False`):
            Whether to return the model outputs along with the loss.
        num_items_in_batch (Optional[torch.Tensor], *optional*):
            The number of items in the batch. If not passed, the loss is computed
            using the default batch size reduction logic.

    Returns:
        The loss of the model along with its output if return_outputs was set to True

    Subclass and override for custom behavior. If you are not using `num_items_in_batch` when computing your loss,
    make sure to overwrite `self.model_accepts_loss_kwargs` to `False`. Otherwise, the loss calculation might be slightly inaccurate when performing gradient accumulation.
    """
    pc = getattr(self.accelerator, "parallelism_config", None)
    if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled and self.model.training:
        return deepspeed_sp_compute_loss(self.accelerator, model, inputs, return_outputs, pc)

    if (self.label_smoother is not None or self.compute_loss_func is not None) and "labels" in inputs:
        labels = inputs.pop("labels")
    else:
        labels = None
    if self.model_accepts_loss_kwargs:
        kwargs = {}
        if num_items_in_batch is not None:
            kwargs["num_items_in_batch"] = num_items_in_batch
        inputs = {**inputs, **kwargs}
    outputs = model(**inputs)

    # User-defined compute_loss function
    if self.compute_loss_func is not None:
        if labels is None:
            logger.warning(
                "Trainer: `compute_loss_func` is defined but `labels=None`. "
                "Your custom loss function will still be called with labels=None. "
            )
        loss = self.compute_loss_func(
            outputs,
            labels,
            num_items_in_batch=num_items_in_batch,
        )
    # Default HF loss handling (label smoothing) if no custom loss function
    elif labels is not None:
        unwrapped_model = self.accelerator.unwrap_model(model)
        model_name = (
            unwrapped_model.base_model.model._get_name()
            if _is_peft_model(unwrapped_model)
            else unwrapped_model._get_name()
        )
        if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
            loss = self.label_smoother(outputs, labels, shift_labels=True)
        else:
            loss = self.label_smoother(outputs, labels)
    else:
        if isinstance(outputs, dict) and "loss" not in outputs:
            raise ValueError(
                "The model did not return a loss from the inputs, only the following keys: "
                f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
            )
        # We don't use .loss here since the model may return tuples instead of ModelOutput.
        loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

    if (
        self.args.average_tokens_across_devices
        and (self.model_accepts_loss_kwargs or self.compute_loss_func)
        and num_items_in_batch is not None
    ):
        loss *= self.accelerator.num_processes if self.args.n_gpu <= 1 else self.args.n_gpu

    return (loss, outputs) if return_outputs else loss

How the loss is computed by Trainer. By default, all models return the loss in the first element.

Args -----= model (nn.Module): The model to compute the loss for. inputs (dict[str, torch.Tensor | Any]): The input data for the model. return_outputs (bool, optional, defaults to False): Whether to return the model outputs along with the loss. num_items_in_batch (Optional[torch.Tensor], optional): The number of items in the batch. If not passed, the loss is computed using the default batch size reduction logic.

Returns -----= The loss of the model along with its output if return_outputs was set to True Subclass and override for custom behavior. If you are not using num_items_in_batch when computing your loss, make sure to overwrite self.model_accepts_loss_kwargs to False. Otherwise, the loss calculation might be slightly inaccurate when performing gradient accumulation.

def compute_loss_context_manager(self) ‑> contextlib.ExitStack

Expand source code

def compute_loss_context_manager(self) -> contextlib.ExitStack:
    """
    A helper wrapper to group together context managers.
    """
    ctx_stack = contextlib.ExitStack()

    autocast_ctx = self.autocast_smart_context_manager()
    if not isinstance(autocast_ctx, contextlib.nullcontext):
        ctx_stack.enter_context(autocast_ctx)

    return ctx_stack

A helper wrapper to group together context managers.

def create_accelerator_and_postprocess(self) ‑> None

Expand source code

def create_accelerator_and_postprocess(self) -> None:
    """Create the accelerator and perform post-creation setup (FSDP, DeepSpeed, etc.)."""
    # We explicitly don't rely on the `Accelerator` to do gradient accumulation
    grad_acc_kwargs = {}
    if self.args.accelerator_config.gradient_accumulation_kwargs is not None:
        grad_acc_kwargs = self.args.accelerator_config.gradient_accumulation_kwargs

    # check if num_steps is attempted to be passed in gradient_accumulation_kwargs
    if "num_steps" in grad_acc_kwargs:
        if self.args.gradient_accumulation_steps > 1:
            # raise because we do not know which setting is intended.
            raise ValueError(
                "The `AcceleratorConfig`'s `num_steps` is set but `gradient_accumulation_steps` is greater than 1 in the passed `TrainingArguments`"
                "If using the passed `AcceleratorConfig` is desired, do not set the `TrainingArguments` `gradient_accumulation_steps`."
            )
        else:
            self.args.gradient_accumulation_steps = grad_acc_kwargs["num_steps"]

    # The Trainer handles GAS itself, so GAS=1 in Accelerate to avoid any double-division
    grad_acc_kwargs["num_steps"] = 1

    # Just making sure that gradient_state have the correct values passed.
    # We don't rely on `accumulate` from accelerate to set sync_gradients in gradient_state.
    # Rather, we do it ourselves by setting self.accelerator.gradient_state._set_sync_gradients.
    gradient_accumulation_plugin = GradientAccumulationPlugin(**grad_acc_kwargs)

    accelerator_config = self.args.accelerator_config.to_dict()

    # Extract dataloader config params from accelerator config
    dataloader_params = ["split_batches", "dispatch_batches", "even_batches", "use_seedable_sampler"]
    dataloader_config = DataLoaderConfiguration(
        **{param: accelerator_config.pop(param) for param in dataloader_params}
    )
    dataloader_config.data_seed = self.args.data_seed

    non_blocking = accelerator_config.pop("non_blocking")

    if non_blocking and not self.args.dataloader_pin_memory:
        logger.warning(
            "`non_blocking` is enabled but `dataloader_pin_memory` is not. For the best performance, it's recommended to enable both."
        )
    dataloader_config.non_blocking = non_blocking
    # this would have been updated above, no need for it anymore
    accelerator_config.pop("gradient_accumulation_kwargs")

    fsdp_plugin = None
    if self.args.fsdp_plugin_args is not None:
        from accelerate.utils import FullyShardedDataParallelPlugin

        fsdp_plugin = FullyShardedDataParallelPlugin(**self.args.fsdp_plugin_args)

    args = self._build_accelerator_args(
        dataloader_config=dataloader_config,
        fsdp_plugin=fsdp_plugin,
        gradient_accumulation_plugin=gradient_accumulation_plugin,
    )

    # create accelerator object
    self.accelerator = Accelerator(**args)
    # some Trainer classes need to use `gather` instead of `gather_for_metrics`, thus we store a flag
    self.gather_function = self.accelerator.gather_for_metrics

    if "use_gather_object" in inspect.signature(self.gather_function).parameters:
        self.gather_function = functools.partial(
            self.gather_function, use_gather_object=self.args.eval_use_gather_object
        )

    # deepspeed and accelerate flags covering both trainer args and accelerate launcher
    self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None
    self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None

    # post accelerator creation setup
    if self.is_fsdp_enabled:
        fsdp_plugin = self.accelerator.state.fsdp_plugin
        for param in ["limit_all_gathers", "activation_checkpointing"]:
            setattr(fsdp_plugin, param, self.args.fsdp_config.get(param, getattr(fsdp_plugin, param)))
        if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing:
            raise ValueError(
                "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg "
                "can't be set to True simultaneously. Please use FSDP's activation_checkpointing logic "
                "when using FSDP."
            )

    if self.is_deepspeed_enabled and getattr(self.args, "hf_deepspeed_config", None) is None:
        propagate_args_to_deepspeed(self.accelerator, self.args)

    # `save_only_model` can't be used with DeepSpeed/FSDP along with `load_best_model_at_end`
    if (
        self.args.save_only_model
        and (self.is_deepspeed_enabled or self.is_fsdp_enabled)
        and self.args.load_best_model_at_end
    ):
        wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP"
        raise ValueError(f"{wrapper} can't be used with `save_only_model` along with `load_best_model_at_end`.")

    # `auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3
    if (
        self.is_deepspeed_enabled
        and self.accelerator.state.deepspeed_plugin.zero_stage == 3
        and self.args.auto_find_batch_size
    ):
        raise ValueError(
            "`auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3. Please consider using Zero-2, Zero-1, or FSDP"
        )
    if (
        self.args.save_only_model
        and self.is_fsdp_enabled
        and "SHARDED_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type)
    ):
        raise ValueError("save_only_model option is not compatible with FSDP state dict type 'SHARDED_STATE_DICT'")

Create the accelerator and perform post-creation setup (FSDP, DeepSpeed, etc.).

def create_model_card(self,
language: str | None = None,
license: str | None = None,
tags: str | list[str] | None = None,
model_name: str | None = None,
finetuned_from: str | None = None,
tasks: str | list[str] | None = None,
dataset_tags: str | list[str] | None = None,
dataset: str | list[str] | None = None,
dataset_args: str | list[str] | None = None) ‑> None

Expand source code

def create_model_card(
    self,
    language: str | None = None,
    license: str | None = None,
    tags: str | list[str] | None = None,
    model_name: str | None = None,
    finetuned_from: str | None = None,
    tasks: str | list[str] | None = None,
    dataset_tags: str | list[str] | None = None,
    dataset: str | list[str] | None = None,
    dataset_args: str | list[str] | None = None,
) -> None:
    """
    Creates a draft of a model card using the information available to the `Trainer`.

    Args:
        language (`str`, *optional*):
            The language of the model (if applicable)
        license (`str`, *optional*):
            The license of the model. Will default to the license of the pretrained model used, if the original
            model given to the `Trainer` comes from a repo on the Hub.
        tags (`str` or `list[str]`, *optional*):
            Some tags to be included in the metadata of the model card.
        model_name (`str`, *optional*):
            The name of the model.
        finetuned_from (`str`, *optional*):
            The name of the model used to fine-tune this one (if applicable). Will default to the name of the repo
            of the original model given to the `Trainer` (if it comes from the Hub).
        tasks (`str` or `list[str]`, *optional*):
            One or several task identifiers, to be included in the metadata of the model card.
        dataset_tags (`str` or `list[str]`, *optional*):
            One or several dataset tags, to be included in the metadata of the model card.
        dataset (`str` or `list[str]`, *optional*):
            One or several dataset identifiers, to be included in the metadata of the model card.
        dataset_args (`str` or `list[str]`, *optional*):
           One or several dataset arguments, to be included in the metadata of the model card.
    """
    if not self.is_world_process_zero():
        return

    model_card_filepath = os.path.join(self.args.output_dir, "README.md")
    is_peft_library = False
    if os.path.exists(model_card_filepath):
        library_name = ModelCard.load(model_card_filepath).data.get("library_name")
        is_peft_library = library_name == "peft"

        # Append existing tags in `tags`
        existing_tags = ModelCard.load(model_card_filepath).data.tags
        if tags is not None and existing_tags is not None:
            if isinstance(tags, str):
                tags = [tags]
            for tag in existing_tags:
                if tag not in tags:
                    tags.append(tag)

    training_summary = TrainingSummary.from_trainer(
        self,
        language=language,
        license=license,
        tags=tags,
        model_name=model_name,
        finetuned_from=finetuned_from,
        tasks=tasks,
        dataset_tags=dataset_tags,
        dataset=dataset,
        dataset_args=dataset_args,
    )
    model_card = training_summary.to_model_card()
    with open(model_card_filepath, "w") as f:
        f.write(model_card)

    if is_peft_library:
        self.accelerator.unwrap_model(self.model).create_or_update_model_card(self.args.output_dir)

Creates a draft of a model card using the information available to the Trainer.

Args -----= language (str, optional): The language of the model (if applicable) license (str, optional): The license of the model. Will default to the license of the pretrained model used, if the original model given to the Trainer comes from a repo on the Hub. tags (str or list[str], optional): Some tags to be included in the metadata of the model card. model_name (str, optional): The name of the model. finetuned_from (str, optional): The name of the model used to fine-tune this one (if applicable). Will default to the name of the repo of the original model given to the Trainer (if it comes from the Hub). tasks (str or list[str], optional): One or several task identifiers, to be included in the metadata of the model card. dataset_tags (str or list[str], optional): One or several dataset tags, to be included in the metadata of the model card. dataset (str or list[str], optional): One or several dataset identifiers, to be included in the metadata of the model card. dataset_args (str or list[str], optional): One or several dataset arguments, to be included in the metadata of the model card.

def create_optimizer(self, model=None) ‑> torch.optim.optimizer.Optimizer

Expand source code

def create_optimizer(self, model=None) -> torch.optim.Optimizer:
    """
    Setup the optimizer.

    We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
    Trainer's init through `optimizers`, or subclass and override this method in a subclass.

    Returns:
        `torch.optim.Optimizer`: The optimizer instance.
    """
    opt_model = self.model if model is None else model

    if self.optimizer is None:
        decay_parameters = self.get_decay_parameter_names(opt_model)
        optimizer_grouped_parameters = [
            {
                "params": [
                    p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                ],
                "weight_decay": self.args.weight_decay,
            },
            {
                "params": [
                    p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                ],
                "weight_decay": 0.0,
            },
        ]

        if self.optimizer_cls_and_kwargs is not None:
            optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
        else:
            optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, opt_model)

        # Check if this is a factory (for complex optimizers like Muon, Dion)
        # Factories are instantiated first, then called with (opt_model, **kwargs)
        if is_optimizer_factory(optimizer_cls):
            self.optimizer = optimizer_cls()(opt_model, **optimizer_kwargs)
        else:
            # Standard optimizer class instantiation
            # Overwrite `params` in case it's created by `get_optimizer_cls_and_kwargs`
            # e.g. for GaLore optimizer.
            if "params" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("params")

            # Overwrite `model` in case it's created by `get_optimizer_cls_and_kwargs`
            # e.g. for LOMO optimizer.
            if "model" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("model")

            # For layer-wise dummy optimizers we overwrite optimizer_grouped_parameters with `optimizer_dict`
            # to avoid arguments conflicts.
            if "optimizer_dict" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("optimizer_dict")

            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

        if "bitsandbytes" in str(optimizer_cls) and optimizer_kwargs.get("optim_bits", None) == 8:
            import bitsandbytes

            manager = bitsandbytes.optim.GlobalOptimManager.get_instance()

            skipped = 0
            for module in opt_model.modules():
                if isinstance(module, nn.Embedding):
                    skipped += sum({p.data_ptr(): p.numel() for p in module.parameters()}.values())
                    logger.info(f"skipped {module}: {skipped / 2**20}M params")
                    manager.register_module_override(module, "weight", {"optim_bits": 32})
                    logger.debug(f"bitsandbytes: will optimize {module} in fp32")
            logger.info(f"skipped: {skipped / 2**20}M params")

    if is_sagemaker_mp_enabled():
        self.optimizer = smp.DistributedOptimizer(self.optimizer)

    return self.optimizer

Setup the optimizer.

We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the Trainer's init through optimizers, or subclass and override this method in a subclass.

Returns -----= torch.optim.Optimizer: The optimizer instance.

def create_optimizer_and_scheduler(self, num_training_steps: int) ‑> None

Expand source code

def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
    """
    Setup the optimizer and the learning rate scheduler.

    We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
    Trainer's init through `optimizers`, or subclass and override this method (or `create_optimizer` and/or
    `create_scheduler`) in a subclass.
    """
    self.create_optimizer()
    self.create_scheduler(num_training_steps=num_training_steps)

Setup the optimizer and the learning rate scheduler.

We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the Trainer's init through optimizers, or subclass and override this method (or create_optimizer and/or create_scheduler) in a subclass.

def create_scheduler(self,
num_training_steps: int,
optimizer: torch.optim.optimizer.Optimizer | None = None) ‑> torch.optim.lr_scheduler.LRScheduler

Expand source code

def create_scheduler(
    self, num_training_steps: int, optimizer: torch.optim.Optimizer | None = None
) -> torch.optim.lr_scheduler.LRScheduler:
    """
    Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or
    passed as an argument.

    Args:
        num_training_steps (int): The number of training steps to do.

    Returns:
        `torch.optim.lr_scheduler.LRScheduler`: The learning rate scheduler instance.
    """
    if self.lr_scheduler is None:
        if optimizer is None:
            if is_sagemaker_mp_enabled() and smp.state.cfg.fp16:
                # If fp16 is enabled, we unwrap the optimizer
                optimizer = self.optimizer.optimizer
            else:
                optimizer = self.optimizer
        self.lr_scheduler = get_scheduler(
            self.args.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
            num_training_steps=num_training_steps,
            scheduler_specific_kwargs=self.args.lr_scheduler_kwargs,
        )
        self._created_lr_scheduler = True
    return self.lr_scheduler

Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or passed as an argument.

Args

-----=

num_training_steps : int

The number of training steps to do.

Returns -----= torch.optim.lr_scheduler.LRScheduler: The learning rate scheduler instance.

def evaluate(self,
eval_dataset: torch.utils.data.dataset.Dataset | dict[str, torch.utils.data.dataset.Dataset] | None = None,
ignore_keys: list[str] | None = None,
metric_key_prefix: str = 'eval') ‑> dict[str, float]

Expand source code

def evaluate(
    self,
    eval_dataset: Dataset | dict[str, Dataset] | None = None,
    ignore_keys: list[str] | None = None,
    metric_key_prefix: str = "eval",
) -> dict[str, float]:
    """
    Run evaluation and returns metrics.

    The calling script will be responsible for providing a method to compute metrics, as they are task-dependent
    (pass it to the init `compute_metrics` argument).

    You can also subclass and override this method to inject custom behavior.

    Args:
        eval_dataset (`Dataset` | dict[str, `Dataset`], *optional*):
            Pass a dataset if you wish to override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns
            not accepted by the `model.forward()` method are automatically removed. If it is a dictionary, it will
            evaluate on each dataset, prepending the dictionary key to the metric name. Datasets must implement the
            `__len__` method.

            <Tip>

            If you pass a dictionary with names of datasets as keys and datasets as values, evaluate will run
            separate evaluations on each dataset. This can be useful to monitor how training affects other
            datasets or simply to get a more fine-grained evaluation.
            When used with `load_best_model_at_end`, make sure `metric_for_best_model` references exactly one
            of the datasets. If you, for example, pass in `{"data1": data1, "data2": data2}` for two datasets
            `data1` and `data2`, you could specify `metric_for_best_model="eval_data1_loss"` for using the
            loss on `data1` and `metric_for_best_model="eval_data2_loss"` for the loss on `data2`.

            </Tip>

        ignore_keys (`list[str]`, *optional*):
            A list of keys in the output of your model (if it is a dictionary) that should be ignored when
            gathering predictions.
        metric_key_prefix (`str`, *optional*, defaults to `"eval"`):
            An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
            "eval_bleu" if the prefix is "eval" (default)

    Returns:
        A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The
        dictionary also contains the epoch number which comes from the training state.
    """
    # handle multiple eval datasets
    override = eval_dataset is not None
    eval_dataset = eval_dataset if override else self.eval_dataset
    if isinstance(eval_dataset, dict):
        metrics = {}
        for eval_dataset_name, _eval_dataset in eval_dataset.items():
            dataset_metrics = self.evaluate(
                eval_dataset=_eval_dataset if override else eval_dataset_name,
                ignore_keys=ignore_keys,
                metric_key_prefix=f"{metric_key_prefix}_{eval_dataset_name}",
            )
            metrics.update(dataset_metrics)
        return metrics

    # memory metrics - must set up as early as possible
    self._memory_tracker.start()

    eval_dataloader = self.get_eval_dataloader(eval_dataset)
    if self.is_fsdp_xla_v2_enabled:
        eval_dataloader = tpu_spmd_dataloader(eval_dataloader)

    start_time = time.time()

    output = self.evaluation_loop(
        eval_dataloader,
        description="Evaluation",
        # No point gathering the predictions if there are no metrics, otherwise we defer to
        # self.args.prediction_loss_only
        prediction_loss_only=True if self.compute_metrics is None else None,
        ignore_keys=ignore_keys,
        metric_key_prefix=metric_key_prefix,
    )

    total_batch_size = self.args.eval_batch_size * self.args.world_size
    if f"{metric_key_prefix}_model_preparation_time" in output.metrics:
        start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"]
    output.metrics.update(
        speed_metrics(
            metric_key_prefix,
            start_time,
            num_samples=output.num_samples,
            num_steps=math.ceil(output.num_samples / total_batch_size),
        )
    )

    self.log(output.metrics)

    if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
        xm.master_print(met.metrics_report())

    self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)

    self._memory_tracker.stop_and_update_metrics(output.metrics)

    return output.metrics

Run evaluation and returns metrics.

The calling script will be responsible for providing a method to compute metrics, as they are task-dependent (pass it to the init compute_metrics argument).

You can also subclass and override this method to inject custom behavior.

Args -----= eval_dataset (Dataset | dict[str, Dataset], optional): Pass a dataset if you wish to override self.eval_dataset. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed. If it is a dictionary, it will evaluate on each dataset, prepending the dictionary key to the metric name. Datasets must implement the __len__ method.

<Tip>

If you pass a dictionary with names of datasets as keys and datasets as values, evaluate will run
separate evaluations on each dataset. This can be useful to monitor how training affects other
datasets or simply to get a more fine-grained evaluation.
When used with <code>load\_best\_model\_at\_end</code>, make sure <code>metric\_for\_best\_model</code> references exactly one
of the datasets. If you, for example, pass in `{"data1": data1, "data2": data2}` for two datasets
<code>data1</code> and <code>data2</code>, you could specify `metric_for_best_model="eval_data1_loss"` for using the
loss on <code>data1</code> and `metric_for_best_model="eval_data2_loss"` for the loss on <code>data2</code>.

</Tip>

ignore_keys (list[str], optional): A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions. metric_key_prefix (str, optional, defaults to "eval"): An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named "eval_bleu" if the prefix is "eval" (default)

Returns -----= A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The dictionary also contains the epoch number which comes from the training state.

def evaluation_loop(self,
dataloader: torch.utils.data.dataloader.DataLoader,
description: str,
prediction_loss_only: bool | None = None,
ignore_keys: list[str] | None = None,
metric_key_prefix: str = 'eval') ‑> transformers.trainer_utils.EvalLoopOutput

Expand source code

def evaluation_loop(
    self,
    dataloader: DataLoader,
    description: str,
    prediction_loss_only: bool | None = None,
    ignore_keys: list[str] | None = None,
    metric_key_prefix: str = "eval",
) -> EvalLoopOutput:
    """
    Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.

    Works both with or without labels.
    """
    args = self.args

    prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only

    # if eval is called w/o train, handle model prep here
    if self.is_deepspeed_enabled and self.deepspeed is None:
        _, _ = deepspeed_init(self, num_training_steps=0, inference=True)

    model = self._wrap_model(self.model, training=False)

    if len(self.accelerator._models) == 0 and model is self.model:
        start_time = time.time()
        model = (
            self.accelerator.prepare(model)
            if self.is_deepspeed_enabled or (self.is_fsdp_enabled and not self.args.torch_compile)
            else self.accelerator.prepare_model(model, evaluation_mode=True)
        )
        self.model_preparation_time = round(time.time() - start_time, 4)

        if self.is_fsdp_enabled:
            self.model = model

        # for the rest of this function `model` is the outside model, whether it was wrapped or not
        if model is not self.model:
            self.model_wrapped = model

        # backward compatibility
        if self.is_deepspeed_enabled:
            self.deepspeed = self.model_wrapped

    # if full fp16 or bf16 eval is wanted and this ``evaluation`` or ``predict`` isn't called
    # while ``train`` is running, cast it to the right dtype first and then put on device
    if not self.is_in_train:
        if args.fp16_full_eval:
            model = model.to(dtype=torch.float16, device=args.device)
        elif args.bf16_full_eval:
            model = model.to(dtype=torch.bfloat16, device=args.device)

    batch_size = self.args.eval_batch_size

    logger.info(f"\n***** Running {description} *****")
    if has_length(dataloader):
        logger.info(f"  Num examples = {self.num_examples(dataloader)}")
    else:
        logger.info("  Num examples: Unknown")
    logger.info(f"  Batch size = {batch_size}")

    if hasattr(model, "eval") and callable(model.eval):
        model.eval()
    if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval):
        self.optimizer.eval()

    self.callback_handler.eval_dataloader = dataloader
    # Do this before wrapping.
    eval_dataset = getattr(dataloader, "dataset", None)

    # Initialize containers
    all_losses = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
    all_preds = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
    all_labels = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
    all_inputs = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)

    metrics = None
    eval_set_kwargs = {}

    # Will be useful when we have an iterable dataset so don't know its length.
    observed_num_examples = 0

    # Main evaluation loop
    for step, inputs in enumerate(dataloader):
        # Update the observed num examples
        observed_batch_size = find_batch_size(inputs)
        if observed_batch_size is not None:
            observed_num_examples += observed_batch_size
            # For batch samplers, batch_size is not known by the dataloader in advance.
            if batch_size is None:
                batch_size = observed_batch_size

        # Prediction step
        losses, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
        main_input_name = getattr(self.model, "main_input_name", "input_ids")
        inputs_decode = (
            self._prepare_input(inputs[main_input_name]) if "inputs" in args.include_for_metrics else None
        )

        if is_torch_xla_available():
            xm.mark_step()

        # Update containers
        if losses is not None:
            losses = self.gather_function(losses.repeat(batch_size))
            all_losses.add(losses)
        if inputs_decode is not None:
            inputs_decode = self.accelerator.pad_across_processes(inputs_decode, dim=1, pad_index=-100)
            inputs_decode = self.gather_function(inputs_decode)
            if not self.args.batch_eval_metrics or description == "Prediction":
                all_inputs.add(inputs_decode)
        if labels is not None:
            # Pad labels here, preparing for preprocess_logits_for_metrics in next logits block.
            labels = self.accelerator.pad_across_processes(labels, dim=1, pad_index=-100)
        if logits is not None:
            logits = self.accelerator.pad_across_processes(logits, dim=1, pad_index=-100)
            if self.preprocess_logits_for_metrics is not None:
                logits = self.preprocess_logits_for_metrics(logits, labels)
            logits = self.gather_function(logits)
            if not self.args.batch_eval_metrics or description == "Prediction":
                all_preds.add(logits)
        if labels is not None:
            labels = self.gather_function(labels)
            if not self.args.batch_eval_metrics or description == "Prediction":
                all_labels.add(labels)

        self.control = self.callback_handler.on_prediction_step(args, self.state, self.control)

        if self.args.batch_eval_metrics:
            if self.compute_metrics is not None and logits is not None and labels is not None:
                is_last_step = self.accelerator.gradient_state.end_of_dataloader
                batch_kwargs = {}
                batch_kwargs["losses"] = losses if "loss" in args.include_for_metrics else None
                batch_kwargs["inputs"] = inputs if "inputs" in args.include_for_metrics else None
                metrics = self.compute_metrics(
                    EvalPrediction(predictions=logits, label_ids=labels, **batch_kwargs),
                    compute_result=is_last_step,
                )

            del losses, logits, labels, inputs
            torch.cuda.empty_cache()

        # Gather all tensors and put them back on the CPU if we have done enough accumulation steps.
        elif args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0:
            all_losses.to_cpu_and_numpy()
            all_preds.to_cpu_and_numpy()
            all_labels.to_cpu_and_numpy()
            all_inputs.to_cpu_and_numpy()

            del losses, logits, labels, inputs
            torch.cuda.empty_cache()

    # After all calls to `.gather_function`, reset to `gather_for_metrics`:
    self.gather_function = self.accelerator.gather_for_metrics

    # Gather all remaining tensors and put them back on the CPU
    all_losses = all_losses.get_arrays()
    all_preds = all_preds.get_arrays()
    all_labels = all_labels.get_arrays()
    all_inputs = all_inputs.get_arrays()

    # Number of samples
    if has_length(eval_dataset):
        num_samples = len(eval_dataset)
    # The instance check is weird and does not actually check for the type, but whether the dataset has the right
    # methods. Therefore we need to make sure it also has the attribute.
    elif isinstance(eval_dataset, IterableDatasetShard) and getattr(eval_dataset, "num_examples", 0) > 0:
        num_samples = eval_dataset.num_examples
    else:
        if has_length(dataloader):
            num_samples = self.num_examples(dataloader)
        else:  # both len(dataloader.dataset) and len(dataloader) fail
            num_samples = observed_num_examples
    if num_samples == 0 and observed_num_examples > 0:
        num_samples = observed_num_examples

    # Metrics!
    if (
        self.compute_metrics is not None
        and all_preds is not None
        and all_labels is not None
        and not self.args.batch_eval_metrics
    ):
        eval_set_kwargs["losses"] = all_losses if "loss" in args.include_for_metrics else None
        eval_set_kwargs["inputs"] = all_inputs if "inputs" in args.include_for_metrics else None
        metrics = self.compute_metrics(
            EvalPrediction(predictions=all_preds, label_ids=all_labels, **eval_set_kwargs)
        )
    elif metrics is None:
        metrics = {}

    # To be JSON-serializable, we need to remove numpy types or zero-d tensors
    metrics = denumpify_detensorize(metrics)

    if isinstance(all_losses, list) and all_losses:
        metrics[f"{metric_key_prefix}_loss"] = np.concatenate(all_losses).mean().item()
    elif isinstance(all_losses, np.ndarray):
        metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()
    if hasattr(self, "model_preparation_time"):
        metrics[f"{metric_key_prefix}_model_preparation_time"] = self.model_preparation_time

    # Prefix all keys with metric_key_prefix + '_'
    for key in list(metrics.keys()):
        if not key.startswith(f"{metric_key_prefix}_"):
            metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

    return EvalLoopOutput(predictions=all_preds, label_ids=all_labels, metrics=metrics, num_samples=num_samples)

Prediction/evaluation loop, shared by Trainer.evaluate() and Trainer.predict().

Works both with or without labels.

def floating_point_ops(self,
inputs: dict[str, torch.Tensor | Any]) ‑> int

Expand source code

def floating_point_ops(self, inputs: dict[str, torch.Tensor | Any]) -> int:
    """
    For models that inherit from [`PreTrainedModel`], uses that method to compute the number of floating point
    operations for every backward + forward pass. If using another model, either implement such a method in the
    model or subclass and override this method.

    Args:
        inputs (`dict[str, torch.Tensor | Any]`):
            The inputs and targets of the model.

    Returns:
        `int`: The number of floating-point operations.
    """
    if (main_input := getattr(self.model, "main_input_name", "input_ids")) in inputs and hasattr(
        self.model, "num_parameters"
    ):
        return 6 * inputs[main_input].numel() * self.model.num_parameters(exclude_embeddings=True)
    return 0

For models that inherit from [PreTrainedModel], uses that method to compute the number of floating point operations for every backward + forward pass. If using another model, either implement such a method in the model or subclass and override this method.

Args -----= inputs (dict[str, torch.Tensor | Any]): The inputs and targets of the model.

Returns -----= int: The number of floating-point operations.

def get_batch_samples(self,
epoch_iterator: ,
num_batches: int,
device: torch.device) ‑> tuple[list, torch.Tensor | int | None]

Expand source code

def get_batch_samples(
    self, epoch_iterator: Iterator, num_batches: int, device: torch.device
) -> tuple[list, torch.Tensor | int | None]:
    """
    Collects a specified number of batches from the epoch iterator and optionally counts the number of items in the batches to properly scale the loss.
    """
    batch_samples = []

    for _ in range(num_batches):
        try:
            batch_samples.append(next(epoch_iterator))
        except StopIteration:
            break

    num_items_in_batch = self._get_num_items_in_batch(batch_samples, device)
    return batch_samples, num_items_in_batch

Collects a specified number of batches from the epoch iterator and optionally counts the number of items in the batches to properly scale the loss.

def get_cp_size(self) ‑> int

Expand source code

def get_cp_size(self) -> int:
    """Get the context parallel size"""
    if getattr(self.accelerator, "parallelism_config", None) is None:
        return 1
    else:
        pc = self.accelerator.parallelism_config
        return pc.cp_size

Get the context parallel size

def get_decay_parameter_names(self, model: torch.nn.modules.module.Module) ‑> list[str]

Expand source code

def get_decay_parameter_names(self, model: nn.Module) -> list[str]:
    """
    Get all parameter names that weight decay will be applied to.

    This function filters out parameters in two ways:
    1. By layer type (instances of layers specified in ALL_LAYERNORM_LAYERS)
    2. By parameter name patterns (containing 'bias', or variation of 'norm')
    """
    forbidden_name_patterns = [r"bias", r"layernorm", r"rmsnorm", r"(?:^|\.)norm(?:$|\.)", r"_norm(?:$|\.)"]
    decay_parameters = get_parameter_names(model, [nn.LayerNorm], forbidden_name_patterns)
    return decay_parameters

Get all parameter names that weight decay will be applied to.

This function filters out parameters in two ways: 1. By layer type (instances of layers specified in ALL_LAYERNORM_LAYERS) 2. By parameter name patterns (containing 'bias', or variation of 'norm')

def get_eval_dataloader(self, eval_dataset: str | torch.utils.data.dataset.Dataset | None = None) ‑> torch.utils.data.dataloader.DataLoader

Expand source code

def get_eval_dataloader(self, eval_dataset: str | Dataset | None = None) -> DataLoader:
    """
    Returns the evaluation [`~torch.utils.data.DataLoader`].

    Subclass and override this method if you want to inject some custom behavior.

    Args:
        eval_dataset (`str` or `torch.utils.data.Dataset`, *optional*):
            If a `str`, will use `self.eval_dataset[eval_dataset]` as the evaluation dataset. If a `Dataset`, will override `self.eval_dataset` and must implement `__len__`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed.
    """
    if eval_dataset is None and self.eval_dataset is None:
        raise ValueError("Trainer: evaluation requires an eval_dataset.")

    # If we have persistent workers, don't do a fork bomb especially as eval datasets
    # don't change during training
    dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval"
    if (
        hasattr(self, "_eval_dataloaders")
        and dataloader_key in self._eval_dataloaders
        and self.args.dataloader_persistent_workers
    ):
        return self._eval_dataloaders[dataloader_key]

    eval_dataset = (
        self.eval_dataset[eval_dataset]
        if isinstance(eval_dataset, str)
        else eval_dataset
        if eval_dataset is not None
        else self.eval_dataset
    )

    return self._get_dataloader(
        dataset=eval_dataset,
        description="Evaluation",
        batch_size=self.args.eval_batch_size,
        sampler_fn=self._get_eval_sampler,
        dataloader_key=dataloader_key,
    )

Returns the evaluation [~torch.utils.data.DataLoader].

Subclass and override this method if you want to inject some custom behavior.

Args -----= eval_dataset (str or torch.utils.data.Dataset, optional): If a str, will use self.eval_dataset[eval_dataset] as the evaluation dataset. If a Dataset, will override self.eval_dataset and must implement __len__. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed.

def get_learning_rates(self) ‑> list[float]

Expand source code

def get_learning_rates(self) -> list[float]:
    """
    Returns the learning rate of each parameter from self.optimizer.
    """
    if self.optimizer is None:
        raise ValueError("Trainer optimizer is None, please make sure you have setup the optimizer before.")
    return [group["lr"] for group in self.optimizer.param_groups]

Returns the learning rate of each parameter from self.optimizer.

def get_num_trainable_parameters(self) ‑> int

Expand source code

def get_num_trainable_parameters(self) -> int:
    """
    Get the number of trainable parameters.
    """
    return sum(p.numel() for p in self.model.parameters() if p.requires_grad)

Get the number of trainable parameters.

def get_optimizer_group(self, param: str | torch.nn.parameter.Parameter | None = None)

Expand source code

def get_optimizer_group(self, param: str | torch.nn.parameter.Parameter | None = None):
    """
    Returns optimizer group for a parameter if given, else returns all optimizer groups for params.

    Args:
        param (`str` or `torch.nn.parameter.Parameter`, *optional*):
            The parameter for which optimizer group needs to be returned.
    """
    if self.optimizer is None:
        raise ValueError("Trainer optimizer is None, please make sure you have setup the optimizer before.")
    if param is not None:
        for group in self.optimizer.param_groups:
            if param in group["params"]:
                return group
    return [group["params"] for group in self.optimizer.param_groups]

Returns optimizer group for a parameter if given, else returns all optimizer groups for params.

Args -----= param (str or torch.nn.parameter.Parameter, optional): The parameter for which optimizer group needs to be returned.

def get_sp_size(self) ‑> int

Expand source code

def get_sp_size(self) -> int:
    """Get the sequence parallel size"""
    if getattr(self.accelerator, "parallelism_config", None) is None:
        return 1
    else:
        pc = self.accelerator.parallelism_config
        return pc.sp_size

Get the sequence parallel size

def get_test_dataloader(self, test_dataset: torch.utils.data.dataset.Dataset) ‑> torch.utils.data.dataloader.DataLoader

Expand source code

def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
    """
    Returns the test [`~torch.utils.data.DataLoader`].

    Subclass and override this method if you want to inject some custom behavior.

    Args:
        test_dataset (`torch.utils.data.Dataset`, *optional*):
            The test dataset to use. If it is a [`~datasets.Dataset`], columns not accepted by the
            `model.forward()` method are automatically removed. It must implement `__len__`.
    """
    return self._get_dataloader(
        dataset=test_dataset,
        description="test",
        batch_size=self.args.eval_batch_size,
        sampler_fn=self._get_eval_sampler,
    )

Returns the test [~torch.utils.data.DataLoader].

Subclass and override this method if you want to inject some custom behavior.

Args -----= test_dataset (torch.utils.data.Dataset, optional): The test dataset to use. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed. It must implement __len__.

def get_total_train_batch_size(self, args: transformers.training_args.TrainingArguments) ‑> int

Expand source code

def get_total_train_batch_size(self, args: TrainingArguments) -> int:
    """Calculates total batch size (micro_batch * grad_accum * dp_world_size).

    Accounts for all parallelism dimensions: TP, CP, and SP.

    Formula: dp_world_size = world_size // (tp_size * cp_size * sp_size)

    Where:
    - TP (Tensor Parallelism): Model layers split across GPUs
    - CP (Context Parallelism): Sequences split using Ring Attention (FSDP2)
    - SP (Sequence Parallelism): Sequences split using ALST/Ulysses (DeepSpeed)

    All dimensions are separate and multiplicative: world_size = dp_size * tp_size * cp_size * sp_size
    """

    dp_world_size = args.world_size // self.get_tp_size() // self.get_cp_size() // self.get_sp_size()
    return self._train_batch_size * args.gradient_accumulation_steps * dp_world_size

Calculates total batch size (micro_batch * grad_accum * dp_world_size).

Accounts for all parallelism dimensions: TP, CP, and SP.

Formula: dp_world_size = world_size // (tp_size * cp_size * sp_size)

Where: - TP (Tensor Parallelism): Model layers split across GPUs - CP (Context Parallelism): Sequences split using Ring Attention (FSDP2) - SP (Sequence Parallelism): Sequences split using ALST/Ulysses (DeepSpeed)

All dimensions are separate and multiplicative: world_size = dp_size * tp_size * cp_size * sp_size

def get_tp_size(self) ‑> int

Expand source code

def get_tp_size(self) -> int:
    """Get the tensor parallel size from either the model or DeepSpeed config."""

    # 1. Check model.tp_size first
    if (model_tp := getattr(self.model, "_tp_size", None)) is not None:
        return model_tp

    # 2. Fall back to DeepSpeed config if enabled
    if self.is_deepspeed_enabled and (deepspeed_config := getattr(self.args, "hf_deepspeed_config", None)):
        return deepspeed_config.config.get("tensor_parallel", {}).get("autotp_size", 1)

    # 3. Default fallback
    return 1

Get the tensor parallel size from either the model or DeepSpeed config.

def get_train_dataloader(self) ‑> torch.utils.data.dataloader.DataLoader

Expand source code

def get_train_dataloader(self) -> DataLoader:
    """
    Returns the training [`~torch.utils.data.DataLoader`].

    Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
    training if necessary) otherwise.

    Subclass and override this method if you want to inject some custom behavior.
    """
    if self.train_dataset is None:
        raise ValueError("Trainer: training requires a train_dataset.")

    return self._get_dataloader(
        dataset=self.train_dataset,
        description="Training",
        batch_size=self._train_batch_size,
        sampler_fn=self._get_train_sampler,
        is_training=True,
    )

Returns the training [~torch.utils.data.DataLoader].

Will use no sampler if train_dataset does not implement __len__, a random sampler (adapted to distributed training if necessary) otherwise.

Subclass and override this method if you want to inject some custom behavior.

def hyperparameter_search(self,
hp_space: Callable[['optuna.Trial'], dict[str, float]] | None = None,
compute_objective: Callable[[dict[str, float]], float] | None = None,
n_trials: int = 20,
direction: str | list[str] = 'minimize',
backend: str | transformers.trainer_utils.HPSearchBackend | None = None,
hp_name: Callable[['optuna.Trial'], str] | None = None,
**kwargs) ‑> transformers.trainer_utils.BestRun | list[transformers.trainer_utils.BestRun]

Expand source code

def hyperparameter_search(
    self,
    hp_space: Callable[["optuna.Trial"], dict[str, float]] | None = None,
    compute_objective: Callable[[dict[str, float]], float] | None = None,
    n_trials: int = 20,
    direction: str | list[str] = "minimize",
    backend: str | HPSearchBackend | None = None,
    hp_name: Callable[["optuna.Trial"], str] | None = None,
    **kwargs,
) -> BestRun | list[BestRun]:
    """
    Launch a hyperparameter search using `optuna` or `Ray Tune`. The optimized quantity is determined
    by `compute_objective`, which defaults to a function returning the evaluation loss when no metric is provided,
    the sum of all metrics otherwise.

    <Tip warning={true}>

    To use this method, you need to have provided a `model_init` when initializing your [`Trainer`]: we need to
    reinitialize the model at each new run. This is incompatible with the `optimizers` argument, so you need to
    subclass [`Trainer`] and override the method [`~Trainer.create_optimizer_and_scheduler`] for custom
    optimizer/scheduler.

    </Tip>

    Args:
        hp_space (`Callable[["optuna.Trial"], dict[str, float]]`, *optional*):
            A function that defines the hyperparameter search space. Will default to
            [`~trainer_utils.default_hp_space_optuna`] or [`~trainer_utils.default_hp_space_ray`]
            depending on your backend.
        compute_objective (`Callable[[dict[str, float]], float]`, *optional*):
            A function computing the objective to minimize or maximize from the metrics returned by the `evaluate`
            method. Will default to [`~trainer_utils.default_compute_objective`].
        n_trials (`int`, *optional*, defaults to 100):
            The number of trial runs to test.
        direction (`str` or `list[str]`, *optional*, defaults to `"minimize"`):
            If it's single objective optimization, direction is `str`, can be `"minimize"` or `"maximize"`, you
            should pick `"minimize"` when optimizing the validation loss, `"maximize"` when optimizing one or
            several metrics. If it's multi objectives optimization, direction is `list[str]`, can be List of
            `"minimize"` and `"maximize"`, you should pick `"minimize"` when optimizing the validation loss,
            `"maximize"` when optimizing one or several metrics.
        backend (`str` or [`~training_utils.HPSearchBackend`], *optional*):
            The backend to use for hyperparameter search. Will default to optuna or Ray Tune, depending
            on which one is installed. If all are installed, will default to optuna.
        hp_name (`Callable[["optuna.Trial"], str]]`, *optional*):
            A function that defines the trial/run name. Will default to None.
        kwargs (`dict[str, Any]`, *optional*):
            Additional keyword arguments for each backend:

            - `optuna`: parameters from
              [optuna.study.create_study](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html)
              and also the parameters `timeout`, `n_jobs` and `gc_after_trial` from
              [optuna.study.Study.optimize](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.optimize)
            - `ray`: parameters from [tune.run](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run).
              If `resources_per_trial` is not set in the `kwargs`, it defaults to 1 CPU core and 1 GPU (if available).
              If `progress_reporter` is not set in the `kwargs`,
              [ray.tune.CLIReporter](https://docs.ray.io/en/latest/tune/api/doc/ray.tune.CLIReporter.html) is used.
    Returns:
        [`trainer_utils.BestRun` or `list[trainer_utils.BestRun]`]: All the information about the best run or best
        runs for multi-objective optimization. Experiment summary can be found in `run_summary` attribute for Ray
        backend.
    """
    if backend is None:
        backend = default_hp_search_backend()
    backend = HPSearchBackend(backend)
    backend_obj = ALL_HYPERPARAMETER_SEARCH_BACKENDS[backend]()
    backend_obj.ensure_available()
    self.hp_search_backend = backend
    if self.model_init is None:
        raise RuntimeError(
            "To use hyperparameter search, you need to pass your model through a model_init function."
        )

    self.hp_space = backend_obj.default_hp_space if hp_space is None else hp_space
    self.hp_name = hp_name
    self.compute_objective = default_compute_objective if compute_objective is None else compute_objective

    best_run = backend_obj.run(self, n_trials, direction, **kwargs)

    self.hp_search_backend = None
    return best_run

Launch a hyperparameter search using optuna or Ray Tune. The optimized quantity is determined by compute_objective, which defaults to a function returning the evaluation loss when no metric is provided, the sum of all metrics otherwise.

To use this method, you need to have provided a model_init when initializing your [Trainer]: we need to reinitialize the model at each new run. This is incompatible with the optimizers argument, so you need to subclass [Trainer] and override the method [~Trainer.create_optimizer_and_scheduler] for custom optimizer/scheduler.

Args -----= hp_space (Callable[["optuna.Trial"], dict[str, float]], optional): A function that defines the hyperparameter search space. Will default to [~trainer_utils.default_hp_space_optuna] or [~trainer_utils.default_hp_space_ray] depending on your backend. compute_objective (Callable[[dict[str, float]], float], optional): A function computing the objective to minimize or maximize from the metrics returned by the evaluate method. Will default to [~trainer_utils.default_compute_objective]. n_trials (int, optional, defaults to 100): The number of trial runs to test. direction (str or list[str], optional, defaults to "minimize"): If it's single objective optimization, direction is str, can be "minimize" or "maximize", you should pick "minimize" when optimizing the validation loss, "maximize" when optimizing one or several metrics. If it's multi objectives optimization, direction is list[str], can be List of "minimize" and "maximize", you should pick "minimize" when optimizing the validation loss, "maximize" when optimizing one or several metrics. backend (str or [~training_utils.HPSearchBackend], optional): The backend to use for hyperparameter search. Will default to optuna or Ray Tune, depending on which one is installed. If all are installed, will default to optuna. hp_name (Callable[["optuna.Trial"], str]], optional): A function that defines the trial/run name. Will default to None. kwargs (dict[str, Any], optional): Additional keyword arguments for each backend:

- <code>optuna</code>: parameters from
  [optuna.study.create_study](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html)
  and also the parameters <code>timeout</code>, <code>n\_jobs</code> and <code>gc\_after\_trial</code> from
  [optuna.study.Study.optimize](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.optimize)
- <code>ray</code>: parameters from [tune.run](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run).
  If <code>resources\_per\_trial</code> is not set in the <code>kwargs</code>, it defaults to 1 CPU core and 1 GPU (if available).
  If <code>progress\_reporter</code> is not set in the <code>kwargs</code>,
  [ray.tune.CLIReporter](https://docs.ray.io/en/latest/tune/api/doc/ray.tune.CLIReporter.html) is used.

Returns -----= [trainer_utils.BestRun or list[trainer_utils.BestRun]]: All the information about the best run or best runs for multi-objective optimization. Experiment summary can be found in run_summary attribute for Ray backend.

def init_hf_repo(self, token: str | None = None) ‑> None

Expand source code

def init_hf_repo(self, token: str | None = None) -> None:
    """
    Initializes a git repo in `self.args.hub_model_id`.
    """
    # Only on process zero
    if not self.is_world_process_zero():
        return

    if self.args.hub_model_id is None:
        repo_name = Path(self.args.output_dir).absolute().name
    else:
        repo_name = self.args.hub_model_id

    token = token if token is not None else self.args.hub_token
    repo_url = create_repo(repo_name, token=token, private=self.args.hub_private_repo, exist_ok=True)
    self.hub_model_id = repo_url.repo_id
    self.push_in_progress = None

Initializes a git repo in self.args.hub_model_id.

def is_local_process_zero(self) ‑> bool

Expand source code

def is_local_process_zero(self) -> bool:
    """
    Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several
    machines) main process.
    """
    return self.args.local_process_index == 0

Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several machines) main process.

def is_world_process_zero(self) ‑> bool

Expand source code

def is_world_process_zero(self) -> bool:
    """
    Whether or not this process is the global main process (when training in a distributed fashion on several
    machines, this is only going to be `True` for one process).
    """
    # Special case for SageMaker ModelParallel since there process_index is dp_process_index, not the global
    # process index.
    if is_sagemaker_mp_enabled():
        return smp.rank() == 0
    return self.args.process_index == 0

Whether or not this process is the global main process (when training in a distributed fashion on several machines, this is only going to be True for one process).

def log(self, logs: dict[str, float], start_time: float | None = None) ‑> None

Expand source code

def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
    """
    Log `logs` on the various objects watching training.

    Subclass and override this method to inject custom behavior.

    Args:
        logs (`dict[str, float]`):
            The values to log.
        start_time (`Optional[float]`):
            The start of training.
    """
    if self.state.epoch is not None:
        logs["epoch"] = self.state.epoch
    if self.args.include_num_input_tokens_seen != "no":
        logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen
        if start_time is not None:
            current_session_num_tokens = self.state.num_input_tokens_seen - self._initial_num_input_tokens_seen
            logs.update(speed_metrics("train", start_time, num_tokens=current_session_num_tokens))

    output = {**logs, "step": self.state.global_step}
    self.state.log_history.append(output)
    self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)

Log logs on the various objects watching training.

Subclass and override this method to inject custom behavior.

Args -----= logs (dict[str, float]): The values to log. start_time (Optional[float]): The start of training.

def log_metrics(self, split, metrics)

Expand source code

def log_metrics(self, split, metrics):
    """
    Log metrics in a specially formatted way.

    Under distributed environment this is done only for a process with rank 0.

    Args:
        split (`str`):
            Mode/split name: one of `train`, `eval`, `test`
        metrics (`dict[str, float]`):
            The metrics returned from train/evaluate/predictmetrics: metrics dict

    Notes on memory reports:

    In order to get memory usage report you need to install `psutil`. You can do that with `pip install psutil`.

    Now when this method is run, you will see a report that will include:

    ```
    init_mem_cpu_alloc_delta   =     1301MB
    init_mem_cpu_peaked_delta  =      154MB
    init_mem_gpu_alloc_delta   =      230MB
    init_mem_gpu_peaked_delta  =        0MB
    train_mem_cpu_alloc_delta  =     1345MB
    train_mem_cpu_peaked_delta =        0MB
    train_mem_gpu_alloc_delta  =      693MB
    train_mem_gpu_peaked_delta =        7MB
    ```

    **Understanding the reports:**

    - the first segment, e.g., `train__`, tells you which stage the metrics are for. Reports starting with `init_`
        will be added to the first stage that gets run. So that if only evaluation is run, the memory usage for the
        `__init__` will be reported along with the `eval_` metrics.
    - the third segment, is either `cpu` or `gpu`, tells you whether it's the general RAM or the gpu0 memory
        metric.
    - `*_alloc_delta` - is the difference in the used/allocated memory counter between the end and the start of the
        stage - it can be negative if a function released more memory than it allocated.
    - `*_peaked_delta` - is any extra memory that was consumed and then freed - relative to the current allocated
        memory counter - it is never negative. When you look at the metrics of any stage you add up `alloc_delta` +
        `peaked_delta` and you know how much memory was needed to complete that stage.

    The reporting happens only for process of rank 0 and gpu 0 (if there is a gpu). Typically this is enough since the
    main process does the bulk of work, but it could be not quite so if model parallel is used and then other GPUs may
    use a different amount of gpu memory. This is also not the same under DataParallel where gpu0 may require much more
    memory than the rest since it stores the gradient and optimizer states for all participating GPUs. Perhaps in the
    future these reports will evolve to measure those too.

    The CPU RAM metric measures RSS (Resident Set Size) includes both the memory which is unique to the process and the
    memory shared with other processes. It is important to note that it does not include swapped out memory, so the
    reports could be imprecise.

    The CPU peak memory is measured using a sampling thread. Due to python's GIL it may miss some of the peak memory if
    that thread didn't get a chance to run when the highest memory was used. Therefore this report can be less than
    reality. Using `tracemalloc` would have reported the exact peak memory, but it doesn't report memory allocations
    outside of python. So if some C++ CUDA extension allocated its own memory it won't be reported. And therefore it
    was dropped in favor of the memory sampling approach, which reads the current process memory usage.

    The GPU allocated and peak memory reporting is done with `torch.cuda.memory_allocated()` and
    `torch.cuda.max_memory_allocated()`. This metric reports only "deltas" for pytorch-specific allocations, as
    `torch.cuda` memory management system doesn't track any memory allocated outside of pytorch. For example, the very
    first cuda call typically loads CUDA kernels, which may take from 0.5 to 2GB of GPU memory.

    Note that this tracker doesn't account for memory allocations outside of [`Trainer`]'s `__init__`, `train`,
    `evaluate` and `predict` calls.

    Because `evaluation` calls may happen during `train`, we can't handle nested invocations because
    `torch.cuda.max_memory_allocated` is a single counter, so if it gets reset by a nested eval call, `train`'s tracker
    will report incorrect info. If this [pytorch issue](https://github.com/pytorch/pytorch/issues/16266) gets resolved
    it will be possible to change this class to be re-entrant. Until then we will only track the outer level of
    `train`, `evaluate` and `predict` methods. Which means that if `eval` is called during `train`, it's the latter
    that will account for its memory usage and that of the former.

    This also means that if any other tool that is used along the [`Trainer`] calls
    `torch.cuda.reset_peak_memory_stats`, the gpu peak memory stats could be invalid. And the [`Trainer`] will disrupt
    the normal behavior of any such tools that rely on calling `torch.cuda.reset_peak_memory_stats` themselves.

    For best performance you may want to consider turning the memory profiling off for production runs.
    """
    if not self.is_world_process_zero():
        return

    print(f"***** {split} metrics *****")
    metrics_formatted = metrics_format(metrics)
    k_width = max(len(str(x)) for x in metrics_formatted)
    v_width = max(len(str(x)) for x in metrics_formatted.values())
    for key in sorted(metrics_formatted.keys()):
        print(f"  {key: <{k_width}} = {metrics_formatted[key]:>{v_width}}")

Log metrics in a specially formatted way.

Under distributed environment this is done only for a process with rank 0.

Args -----= split (str): Mode/split name: one of train, eval, test metrics (dict[str, float]): The metrics returned from train/evaluate/predictmetrics: metrics dict Notes on memory reports:

In order to get memory usage report you need to install psutil. You can do that with pip install psutil.

Now when this method is run, you will see a report that will include:

init_mem_cpu_alloc_delta   =     1301MB
init_mem_cpu_peaked_delta  =      154MB
init_mem_gpu_alloc_delta   =      230MB
init_mem_gpu_peaked_delta  =        0MB
train_mem_cpu_alloc_delta  =     1345MB
train_mem_cpu_peaked_delta =        0MB
train_mem_gpu_alloc_delta  =      693MB
train_mem_gpu_peaked_delta =        7MB

Understanding the reports:

• the first segment, e.g., train__, tells you which stage the metrics are for. Reports starting with init_ will be added to the first stage that gets run. So that if only evaluation is run, the memory usage for the __init__ will be reported along with the eval_ metrics.
• the third segment, is either cpu or gpu, tells you whether it's the general RAM or the gpu0 memory metric.
• *_alloc_delta - is the difference in the used/allocated memory counter between the end and the start of the stage - it can be negative if a function released more memory than it allocated.
• *_peaked_delta - is any extra memory that was consumed and then freed - relative to the current allocated memory counter - it is never negative. When you look at the metrics of any stage you add up alloc_delta + peaked_delta and you know how much memory was needed to complete that stage.

The reporting happens only for process of rank 0 and gpu 0 (if there is a gpu). Typically this is enough since the main process does the bulk of work, but it could be not quite so if model parallel is used and then other GPUs may use a different amount of gpu memory. This is also not the same under DataParallel where gpu0 may require much more memory than the rest since it stores the gradient and optimizer states for all participating GPUs. Perhaps in the future these reports will evolve to measure those too.

The CPU RAM metric measures RSS (Resident Set Size) includes both the memory which is unique to the process and the memory shared with other processes. It is important to note that it does not include swapped out memory, so the reports could be imprecise.

The CPU peak memory is measured using a sampling thread. Due to python's GIL it may miss some of the peak memory if that thread didn't get a chance to run when the highest memory was used. Therefore this report can be less than reality. Using tracemalloc would have reported the exact peak memory, but it doesn't report memory allocations outside of python. So if some C++ CUDA extension allocated its own memory it won't be reported. And therefore it was dropped in favor of the memory sampling approach, which reads the current process memory usage.

The GPU allocated and peak memory reporting is done with torch.cuda.memory_allocated() and torch.cuda.max_memory_allocated(). This metric reports only "deltas" for pytorch-specific allocations, as torch.cuda memory management system doesn't track any memory allocated outside of pytorch. For example, the very first cuda call typically loads CUDA kernels, which may take from 0.5 to 2GB of GPU memory.

Note that this tracker doesn't account for memory allocations outside of [Trainer]'s __init__, train, evaluate and predict calls.

Because evaluation calls may happen during train, we can't handle nested invocations because torch.cuda.max_memory_allocated is a single counter, so if it gets reset by a nested eval call, train's tracker will report incorrect info. If this pytorch issue gets resolved it will be possible to change this class to be re-entrant. Until then we will only track the outer level of train, evaluate and predict methods. Which means that if eval is called during train, it's the latter that will account for its memory usage and that of the former.

This also means that if any other tool that is used along the [Trainer] calls torch.cuda.reset_peak_memory_stats, the gpu peak memory stats could be invalid. And the [Trainer] will disrupt the normal behavior of any such tools that rely on calling torch.cuda.reset_peak_memory_stats themselves.

For best performance you may want to consider turning the memory profiling off for production runs.

def metrics_format(metrics: dict[str, float]) ‑> dict[str, float]

Expand source code

def metrics_format(metrics: dict[str, float]) -> dict[str, float]:
    """
    Reformat Trainer metrics values to a human-readable format.

    Args:
        metrics (`dict[str, float]`):
            The metrics returned from train/evaluate/predict

    Returns:
        metrics (`dict[str, float]`): The reformatted metrics
    """

    metrics_copy = metrics.copy()
    for k, v in metrics_copy.items():
        if "_mem_" in k:
            metrics_copy[k] = f"{v >> 20}MB"
        elif "_runtime" in k:
            metrics_copy[k] = _secs2timedelta(v)
        elif k == "total_flos":
            metrics_copy[k] = f"{int(v) >> 30}GF"
        elif isinstance(metrics_copy[k], float):
            metrics_copy[k] = round(v, 4)

    return metrics_copy

Reformat Trainer metrics values to a human-readable format.

Args -----= metrics (dict[str, float]): The metrics returned from train/evaluate/predict

Returns -----= metrics (dict[str, float]): The reformatted metrics

def num_examples(self, dataloader: torch.utils.data.dataloader.DataLoader) ‑> int

Expand source code

def num_examples(self, dataloader: DataLoader) -> int:
    """
    Helper to get number of samples in a [`~torch.utils.data.DataLoader`] by accessing its dataset. When
    dataloader.dataset does not exist or has no length, estimates as best it can
    """
    try:
        dataset = dataloader.dataset
        # Special case for IterableDatasetShard, we need to dig deeper
        if isinstance(dataset, IterableDatasetShard):
            return len(dataloader.dataset.dataset)
        return len(dataloader.dataset)
    except (NameError, AttributeError, TypeError):  # no dataset or length, estimate by length of dataloader
        return len(dataloader) * self.args.per_device_train_batch_size

Helper to get number of samples in a [~torch.utils.data.DataLoader] by accessing its dataset. When dataloader.dataset does not exist or has no length, estimates as best it can

def pop_callback(self,
callback: type[transformers.trainer_callback.TrainerCallback] | transformers.trainer_callback.TrainerCallback) ‑> transformers.trainer_callback.TrainerCallback | None

Expand source code

def pop_callback(self, callback: type[TrainerCallback] | TrainerCallback) -> TrainerCallback | None:
    """
    Remove a callback from the current list of [`~transformers.TrainerCallback`] and returns it.

    If the callback is not found, returns `None` (and no error is raised).

    Args:
       callback (`type` or [`~transformers.TrainerCallback]`):
           A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
           first case, will pop the first member of that class found in the list of callbacks.

    Returns:
        [`~transformers.TrainerCallback`]: The callback removed, if found.
    """
    return self.callback_handler.pop_callback(callback)

Remove a callback from the current list of [~transformers.TrainerCallback] and returns it.

If the callback is not found, returns None (and no error is raised).

Args -----= callback (type or [~transformers.TrainerCallback]): A [~transformers.TrainerCallback] class or an instance of a [~transformers.TrainerCallback]. In the first case, will pop the first member of that class found in the list of callbacks.

Returns -----= [~transformers.TrainerCallback]: The callback removed, if found.

def predict(self,
test_dataset: torch.utils.data.dataset.Dataset,
ignore_keys: list[str] | None = None,
metric_key_prefix: str = 'test') ‑> transformers.trainer_utils.PredictionOutput

Expand source code

def predict(
    self, test_dataset: Dataset, ignore_keys: list[str] | None = None, metric_key_prefix: str = "test"
) -> PredictionOutput:
    """
    Run prediction and returns predictions and potential metrics.

    Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method
    will also return metrics, like in `evaluate()`.

    Args:
        test_dataset (`Dataset`):
            Dataset to run the predictions on. If it is an `datasets.Dataset`, columns not accepted by the
            `model.forward()` method are automatically removed. Has to implement the method `__len__`
        ignore_keys (`list[str]`, *optional*):
            A list of keys in the output of your model (if it is a dictionary) that should be ignored when
            gathering predictions.
        metric_key_prefix (`str`, *optional*, defaults to `"test"`):
            An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
            "test_bleu" if the prefix is "test" (default)

    <Tip>

    If your predictions or labels have different sequence length (for instance because you're doing dynamic padding
    in a token classification task) the predictions will be padded (on the right) to allow for concatenation into
    one array. The padding index is -100.

    </Tip>

    Returns: *NamedTuple* A namedtuple with the following keys:

        - predictions (`np.ndarray`): The predictions on `test_dataset`.
        - label_ids (`np.ndarray`, *optional*): The labels (if the dataset contained some).
        - metrics (`dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained
          labels).
    """
    # memory metrics - must set up as early as possible
    self._memory_tracker.start()

    test_dataloader = self.get_test_dataloader(test_dataset)
    start_time = time.time()

    output = self.evaluation_loop(
        test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
    )
    total_batch_size = self.args.eval_batch_size * self.args.world_size
    if f"{metric_key_prefix}_model_preparation_time" in output.metrics:
        start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"]
    output.metrics.update(
        speed_metrics(
            metric_key_prefix,
            start_time,
            num_samples=output.num_samples,
            num_steps=math.ceil(output.num_samples / total_batch_size),
        )
    )

    self.control = self.callback_handler.on_predict(self.args, self.state, self.control, output.metrics)
    self._memory_tracker.stop_and_update_metrics(output.metrics)

    return PredictionOutput(predictions=output.predictions, label_ids=output.label_ids, metrics=output.metrics)

Run prediction and returns predictions and potential metrics.

Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method will also return metrics, like in evaluate().

Args -----= test_dataset (Dataset): Dataset to run the predictions on. If it is an datasets.Dataset, columns not accepted by the model.forward() method are automatically removed. Has to implement the method __len__ ignore_keys (list[str], optional): A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions. metric_key_prefix (str, optional, defaults to "test"): An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named "test_bleu" if the prefix is "test" (default)

If your predictions or labels have different sequence length (for instance because you're doing dynamic padding in a token classification task) the predictions will be padded (on the right) to allow for concatenation into one array. The padding index is -100.

Returns: NamedTuple A namedtuple with the following keys:

- predictions (<code>np.ndarray</code>): The predictions on <code>test\_dataset</code>.
- label_ids (<code>np.ndarray</code>, *optional*): The labels (if the dataset contained some).
- metrics (<code>dict\[str, float]</code>, *optional*): The potential dictionary of metrics (if the dataset contained
  labels).

def prediction_step(self,
model: torch.nn.modules.module.Module,
inputs: dict[str, torch.Tensor | Any],
prediction_loss_only: bool,
ignore_keys: list[str] | None = None) ‑> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]

Expand source code

def prediction_step(
    self,
    model: nn.Module,
    inputs: dict[str, torch.Tensor | Any],
    prediction_loss_only: bool,
    ignore_keys: list[str] | None = None,
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
    """
    Perform an evaluation step on `model` using `inputs`.

    Subclass and override to inject custom behavior.

    Args:
        model (`nn.Module`):
            The model to evaluate.
        inputs (`dict[str, torch.Tensor | Any]`):
            The inputs and targets of the model.

            The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
            argument `labels`. Check your model's documentation for all accepted arguments.
        prediction_loss_only (`bool`):
            Whether or not to return the loss only.
        ignore_keys (`list[str]`, *optional*):
            A list of keys in the output of your model (if it is a dictionary) that should be ignored when
            gathering predictions.

    Return:
        tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss,
        logits and labels (each being optional).
    """
    has_labels = False if len(self.label_names) == 0 else all(inputs.get(k) is not None for k in self.label_names)
    # For CLIP-like models capable of returning loss values.
    # If `return_loss` is not specified or being `None` in `inputs`, we check if the default value of `return_loss`
    # is `True` in `model.forward`.
    return_loss = inputs.get("return_loss")
    if return_loss is None:
        return_loss = self.can_return_loss
    loss_without_labels = len(self.label_names) == 0 and return_loss

    inputs = self._prepare_inputs(inputs)
    if ignore_keys is None:
        if hasattr(self.model, "config"):
            ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", ["past_key_values"])
        else:
            ignore_keys = []

    # labels may be popped when computing the loss (label smoothing for instance) so we grab them first.
    if has_labels or loss_without_labels:
        labels = nested_detach(tuple(inputs.get(name) for name in self.label_names))
        if len(labels) == 1:
            labels = labels[0]
    else:
        labels = None

    with torch.no_grad():
        if is_sagemaker_mp_enabled():
            raw_outputs = smp_forward_only(model, inputs)
            if has_labels or loss_without_labels:
                if isinstance(raw_outputs, dict):
                    loss_mb = raw_outputs["loss"]
                    logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys + ["loss"])
                else:
                    loss_mb = raw_outputs[0]
                    logits_mb = raw_outputs[1:]

                loss = loss_mb.reduce_mean().detach().cpu()
                logits = smp_nested_concat(logits_mb)
            else:
                loss = None
                if isinstance(raw_outputs, dict):
                    logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys)
                else:
                    logits_mb = raw_outputs
                logits = smp_nested_concat(logits_mb)
        else:
            if has_labels or loss_without_labels:
                with self.compute_loss_context_manager():
                    num_items_in_batch = self._get_num_items_in_batch([inputs], self.args.device)
                    loss, outputs = self.compute_loss(
                        model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch
                    )
                loss = loss.detach().mean()

                if isinstance(outputs, dict):
                    logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
                else:
                    logits = outputs[1:]
            else:
                loss = None
                with self.compute_loss_context_manager():
                    outputs = model(**inputs)
                if isinstance(outputs, dict):
                    logits = tuple(v for k, v in outputs.items() if k not in ignore_keys)
                else:
                    logits = outputs

    if prediction_loss_only:
        return (loss, None, None)

    logits = nested_detach(logits)
    if len(logits) == 1:
        logits = logits[0]

    return (loss, logits, labels)

Perform an evaluation step on model using inputs.

Subclass and override to inject custom behavior.

Args -----= model (nn.Module): The model to evaluate. inputs (dict[str, torch.Tensor | Any]): The inputs and targets of the model.

The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
argument <code>labels</code>. Check your model's documentation for all accepted arguments.

prediction_loss_only (bool): Whether or not to return the loss only. ignore_keys (list[str], optional): A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions.

Return -----= tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and labels (each being optional).

def push_to_hub(self,
commit_message: str | None = 'End of training',
blocking: bool = True,
token: str | None = None,
revision: str | None = None,
**kwargs) ‑> huggingface_hub.hf_api.CommitInfo

Expand source code

def push_to_hub(
    self,
    commit_message: str | None = "End of training",
    blocking: bool = True,
    token: str | None = None,
    revision: str | None = None,
    **kwargs,
) -> CommitInfo:
    """
    Upload `self.model` and `self.processing_class` to the 🤗 model hub on the repo `self.args.hub_model_id`.

    Parameters:
        commit_message (`str`, *optional*, defaults to `"End of training"`):
            Message to commit while pushing.
        blocking (`bool`, *optional*, defaults to `True`):
            Whether the function should return only when the `git push` has finished.
        token (`str`, *optional*, defaults to `None`):
            Token with write permission to overwrite Trainer's original args.
        revision (`str`, *optional*):
            The git revision to commit from. Defaults to the head of the "main" branch.
        kwargs (`dict[str, Any]`, *optional*):
            Additional keyword arguments passed along to [`~Trainer.create_model_card`].

    Returns:
        The URL of the repository where the model was pushed if `blocking=False`, or a `Future` object tracking the
        progress of the commit if `blocking=True`.
    """
    self.callback_handler.on_push_begin(self.args, self.state, self.control)

    model_name = kwargs.pop("model_name", None)
    if model_name is None and self.args.should_save:
        if self.args.hub_model_id is None:
            model_name = Path(self.args.output_dir).name
        else:
            model_name = self.args.hub_model_id.split("/")[-1]
    token = token if token is not None else self.args.hub_token

    # In case the user calls this method with args.push_to_hub = False
    if self.hub_model_id is None:
        self.init_hf_repo(token=token)

    # Needs to be executed on all processes for TPU training, but will only save on the processed determined by
    # self.args.should_save.
    self.save_model(_internal_call=True)

    # Only push from one node.
    if not self.is_world_process_zero():
        return

    # Add additional tags in the case the model has already some tags and users pass
    # "tags" argument to `push_to_hub` so that trainer automatically handles internal tags
    # from all models since Trainer does not call `model.push_to_hub`.
    if getattr(self.model, "model_tags", None) is not None:
        if "tags" not in kwargs:
            kwargs["tags"] = []

        # If it is a string, convert it to a list
        if isinstance(kwargs["tags"], str):
            kwargs["tags"] = [kwargs["tags"]]

        for model_tag in self.model.model_tags:
            if model_tag not in kwargs["tags"]:
                kwargs["tags"].append(model_tag)

    self.create_model_card(model_name=model_name, **kwargs)

    if revision is None:
        revision = self.args.hub_revision

    # Wait for the current upload to be finished.
    self._finish_current_push()

    return upload_folder(
        repo_id=self.hub_model_id,
        folder_path=self.args.output_dir,
        commit_message=commit_message,
        token=token,
        run_as_future=not blocking,
        ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"],
        revision=revision,
    )

Upload self.model and self.processing_class to the 🤗 model hub on the repo self.args.hub_model_id.

Parameters -----= commit_message (str, optional, defaults to "End of training"): Message to commit while pushing. blocking (bool, optional, defaults to True): Whether the function should return only when the git push has finished. token (str, optional, defaults to None): Token with write permission to overwrite Trainer's original args. revision (str, optional): The git revision to commit from. Defaults to the head of the "main" branch. kwargs (dict[str, Any], optional): Additional keyword arguments passed along to [~Trainer.create_model_card].

Returns -----= The URL of the repository where the model was pushed if blocking=False, or a Future object tracking the progress of the commit if blocking=True.

def remove_callback(self,
callback: type[transformers.trainer_callback.TrainerCallback] | transformers.trainer_callback.TrainerCallback) ‑> None

Expand source code

def remove_callback(self, callback: type[TrainerCallback] | TrainerCallback) -> None:
    """
    Remove a callback from the current list of [`~transformers.TrainerCallback`].

    Args:
       callback (`type` or [`~transformers.TrainerCallback]`):
           A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
           first case, will remove the first member of that class found in the list of callbacks.
    """
    self.callback_handler.remove_callback(callback)

Remove a callback from the current list of [~transformers.TrainerCallback].

Args -----= callback (type or [~transformers.TrainerCallback]): A [~transformers.TrainerCallback] class or an instance of a [~transformers.TrainerCallback]. In the first case, will remove the first member of that class found in the list of callbacks.

def save_metrics(self, split, metrics, combined=True)

Expand source code

def save_metrics(self, split, metrics, combined=True):
    """
    Save metrics into a json file for that split, e.g. `train_results.json`.

    Under distributed environment this is done only for a process with rank 0.

    Args:
        split (`str`):
            Mode/split name: one of `train`, `eval`, `test`, `all`
        metrics (`dict[str, float]`):
            The metrics returned from train/evaluate/predict
        combined (`bool`, *optional*, defaults to `True`):
            Creates combined metrics by updating `all_results.json` with metrics of this call

    To understand the metrics please read the docstring of [`~Trainer.log_metrics`]. The only difference is that raw
    unformatted numbers are saved in the current method.

    """
    if not self.is_world_process_zero():
        return

    path = os.path.join(self.args.output_dir, f"{split}_results.json")
    with open(path, "w") as f:
        json.dump(metrics, f, indent=4, sort_keys=True)

    if combined:
        path = os.path.join(self.args.output_dir, "all_results.json")
        if os.path.exists(path):
            with open(path) as f:
                all_metrics = json.load(f)
        else:
            all_metrics = {}

        all_metrics.update(metrics)
        with open(path, "w") as f:
            json.dump(all_metrics, f, indent=4, sort_keys=True)

Save metrics into a json file for that split, e.g. train_results.json.

Under distributed environment this is done only for a process with rank 0.

Args -----= split (str): Mode/split name: one of train, eval, test, all metrics (dict[str, float]): The metrics returned from train/evaluate/predict combined (bool, optional, defaults to True): Creates combined metrics by updating all_results.json with metrics of this call To understand the metrics please read the docstring of [~Trainer.log_metrics]. The only difference is that raw unformatted numbers are saved in the current method.

def save_model(self, output_dir: str | None = None) ‑> None

Expand source code

def save_model(self, output_dir: str | None = None, _internal_call: bool = False) -> None:
    """
    Will save the model, so you can reload it using `from_pretrained()`.

    Will only save from the main process.
    """

    if output_dir is None:
        output_dir = self.args.output_dir

    if is_torch_xla_available():
        save_tpu_checkpoint(
            self.model, self.args, self.accelerator, self.processing_class, self.is_fsdp_xla_v1_enabled, output_dir
        )
    elif is_sagemaker_mp_enabled():
        # Calling the state_dict needs to be done on the wrapped model and on all processes.
        os.makedirs(output_dir, exist_ok=True)
        state_dict = self.model_wrapped.state_dict()
        if self.args.should_save:
            self._save(output_dir, state_dict=state_dict)
        Path(os.path.join(output_dir, "user_content.pt")).touch()
    elif self.is_fsdp_enabled:
        if "FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type):
            state_dict = self.accelerator.get_state_dict(self.model)
            if self.args.should_save:
                self._save(output_dir, state_dict=state_dict)
    elif self.is_deepspeed_enabled:
        try:
            accept_exclude_frozen_parameters = "exclude_frozen_parameters" in set(
                inspect.signature(self.model_wrapped.save_checkpoint).parameters.keys()
            )
            zero3_sharding = self.deepspeed.config.get("zero_optimization", {}).get("stage", None) == 3
            if accept_exclude_frozen_parameters and _is_peft_model(self.model) and zero3_sharding:
                # When using PEFT with DeepSpeed ZeRO Stage 3,
                # we do not need to load the frozen parameters
                state_dict = self.deepspeed._zero3_consolidated_16bit_state_dict(exclude_frozen_parameters=True)
            else:
                state_dict = self.accelerator.get_state_dict(self.deepspeed)
            if self.args.should_save:
                self._save(output_dir, state_dict=state_dict)
        except ValueError:
            logger.warning(
                " stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead, use"
                " zero_to_fp32.py to recover weights"
            )
            if self.args.should_save:
                self._save(output_dir, state_dict={})
            # remove the dummy state_dict
            remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
            self.model_wrapped.save_checkpoint(output_dir)

    elif self.args.should_save:
        self._save(output_dir)

    # Push to the Hub when `save_model` is called by the user.
    if self.args.push_to_hub and not _internal_call:
        self.push_to_hub(commit_message="Model save", revision=self.args.hub_revision)

Will save the model, so you can reload it using from_pretrained().

Will only save from the main process.

def save_state(self)

Expand source code

def save_state(self):
    """
    Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model.

    Under distributed environment this is done only for a process with rank 0.
    """
    if not self.is_world_process_zero():
        return

    path = os.path.join(self.args.output_dir, "trainer_state.json")
    self.state.save_to_json(path)

Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model.

Under distributed environment this is done only for a process with rank 0.

def set_initial_training_values(self,
args: transformers.training_args.TrainingArguments,
dataloader: torch.utils.data.dataloader.DataLoader) ‑> tuple[int, int, int, int, int, int | None, int]

Expand source code

def set_initial_training_values(
    self, args: TrainingArguments, dataloader: DataLoader
) -> tuple[int, int, int, int, int, int | None, int]:
    """
    Calculates and returns the following values:
    - `num_train_epochs`
    - `num_update_steps_per_epoch`
    - `num_examples`
    - `num_train_samples`
    - `total_train_batch_size`
    - `steps_in_epoch` (total batches per epoch)
    - `max_steps`
    """
    # Case 1: we rely on `args.max_steps` first
    max_steps = args.max_steps
    # If max_steps is negative, we use the number of epochs to determine the number of total steps later
    epoch_based = max_steps < 0
    len_dataloader = len(dataloader) if has_length(dataloader) else None
    total_train_batch_size = self.get_total_train_batch_size(args)

    # Account for Sequence Parallelism (SP) dataloader adapter's effect
    sp_size = self.get_sp_size()
    if sp_size > 1 and len_dataloader is not None:
        len_dataloader = len_dataloader * sp_size

    # Case 2: We have a dataloader length and can extrapolate
    if len_dataloader is not None:
        num_update_steps_per_epoch = max(
            len_dataloader // args.gradient_accumulation_steps
            + int(len_dataloader % args.gradient_accumulation_steps > 0),
            1,
        )
        # Case 3: We have a length but are using epochs, we can extrapolate the number of steps
        if epoch_based:
            max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
    # Now we figure out `num_examples`, `num_train_epochs`, and `train_samples`
    if len_dataloader:
        num_examples = self.num_examples(dataloader)
        if args.max_steps > 0:
            num_train_epochs = max_steps // num_update_steps_per_epoch + int(
                max_steps % num_update_steps_per_epoch > 0
            )
            # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's
            # the best we can do.
            num_train_samples = max_steps * total_train_batch_size
        else:
            num_train_epochs = math.ceil(args.num_train_epochs)
            num_train_samples = self.num_examples(dataloader) * args.num_train_epochs
    elif args.max_steps > 0:  # Rely on max_steps when dataloader does not have a working size
        # Setting a very large number of epochs so we go as many times as necessary over the iterator.
        num_train_epochs = sys.maxsize
        num_update_steps_per_epoch = max_steps
        num_examples = total_train_batch_size * args.max_steps
        num_train_samples = args.max_steps * total_train_batch_size
    else:
        raise ValueError(
            "args.max_steps must be set to a positive value if dataloader does not have a length, was"
            f" {args.max_steps}"
        )
    steps_in_epoch = len_dataloader if len_dataloader is not None else max_steps * args.gradient_accumulation_steps
    return (
        num_train_epochs,
        num_update_steps_per_epoch,
        num_examples,
        num_train_samples,
        total_train_batch_size,
        steps_in_epoch,
        max_steps,
    )

Calculates and returns the following values: - num_train_epochs - num_update_steps_per_epoch - num_examples - num_train_samples - total_train_batch_size - steps_in_epoch (total batches per epoch) - max_steps

def store_flos(self) ‑> None

Expand source code

def store_flos(self) -> None:
    """Store the number of floating-point operations that went into the model."""
    if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
        self.state.total_flos += (
            distributed_broadcast_scalars([self.current_flos], device=self.args.device).sum().item()
        )
        self.current_flos = 0
    else:
        self.state.total_flos += self.current_flos
        self.current_flos = 0

Store the number of floating-point operations that went into the model.

def train(self,
resume_from_checkpoint: str | bool | None = None,
trial: optuna.Trial | dict[str, Any] | None = None,
ignore_keys_for_eval: list[str] | None = None) ‑> transformers.trainer_utils.TrainOutput

Expand source code

def train(
    self,
    resume_from_checkpoint: str | bool | None = None,
    trial: "optuna.Trial | dict[str, Any] | None" = None,
    ignore_keys_for_eval: list[str] | None = None,
) -> TrainOutput:
    """
    Main training entry point.

    Args:
        resume_from_checkpoint (`str` or `bool`, *optional*):
            If a `str`, local path to a saved checkpoint as saved by a previous instance of [`Trainer`]. If a
            `bool` and equals `True`, load the last checkpoint in *args.output_dir* as saved by a previous instance
            of [`Trainer`]. If present, training will resume from the model/optimizer/scheduler states loaded here.
        trial (`optuna.Trial` or `dict[str, Any]`, *optional*):
            The trial run or the hyperparameter dictionary for hyperparameter search.
        ignore_keys_for_eval (`list[str]`, *optional*)
            A list of keys in the output of your model (if it is a dictionary) that should be ignored when
            gathering predictions for evaluation during the training.

    Returns:
        [`~trainer_utils.TrainOutput`]: Object containing the global step count, training loss, and metrics.
    """
    if resume_from_checkpoint is False:
        resume_from_checkpoint = None

    # memory metrics - must set up as early as possible
    self._memory_tracker.start()

    args = self.args

    self.is_in_train = True

    # Model re-init
    if self.model_init is not None:
        # Seed must be set before instantiating the model when using model_init.
        enable_full_determinism(args.seed) if args.full_determinism else set_seed(args.seed)
        self.model = self.call_model_init(trial)
        # Reinitializes optimizer and scheduler
        self.optimizer, self.lr_scheduler = None, None
        if self.place_model_on_device:
            self._move_model_to_device(self.model, args.device)
        self.model_wrapped = self.model

    if self.args.use_liger_kernel:
        apply_liger_kernel(self.model, self.args.liger_kernel_config)

    # When fp16/bf16 full eval is enabled, __init__ skips device placement so that
    # evaluation_loop can cast dtype and move in one step. Move the model now for training.
    if (args.fp16_full_eval or args.bf16_full_eval) and not self.is_model_parallel and self.model_init is None:
        self._move_model_to_device(self.model, args.device)

    # Activate gradient checkpointing if needed
    if args.gradient_checkpointing:
        self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=args.gradient_checkpointing_kwargs)

    # If the model uses a tokenizer, it may have a new tokens for fine-tuning purposes.
    if isinstance(self.processing_class, (PreTrainedTokenizerBase, ProcessorMixin)) and hasattr(
        self.model, "config"
    ):
        align_special_tokens(self.model, self.processing_class)

    # Attach NEFTune hooks if necessary
    if self.neftune_noise_alpha is not None:
        self.neftune_hook_handle = activate_neftune(self.model, self.neftune_noise_alpha, self.accelerator)

    # This might change the seed so needs to run first.
    self._hp_search_setup(trial)

    if DebugOption.UNDERFLOW_OVERFLOW in args.debug:
        if args.n_gpu > 1:
            # nn.DataParallel(model) replicates the model, creating new variables and module
            # references registered here no longer work on other gpus, breaking the module
            raise ValueError(
                "Currently --debug underflow_overflow is not supported under DP. Please use DDP with torchrun"
            )
        else:
            DebugUnderflowOverflow(self.model)

    # Load potential model checkpoint
    if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint:
        resume_from_checkpoint = get_last_checkpoint(args.output_dir)
        if resume_from_checkpoint is None:
            raise ValueError(f"No valid checkpoint found in output directory ({args.output_dir})")

    if resume_from_checkpoint is not None:
        # Load model checkpoint before accelerator.prepare() for regular models,
        # so that buffers and parameters are on the right device after prepare.
        # Deepspeed/FSDP models are loaded after prepare in _prepare_for_training.
        if not is_sagemaker_mp_enabled() and not self.is_deepspeed_enabled and not self.is_fsdp_enabled:
            self._load_from_checkpoint(resume_from_checkpoint)
        state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
        if state.train_batch_size is not None and args.auto_find_batch_size:
            # Only restore the checkpoint's train_batch_size when using auto_find_batch_size,
            self._train_batch_size = state.train_batch_size

    inner_training_loop = find_executable_batch_size(
        self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size
    )
    # Disable progress bars when uploading models during checkpoints to avoid polluting stdout
    ctx = suppress_progress_bars() if args.push_to_hub else contextlib.nullcontext()
    with ctx:
        return inner_training_loop(
            args=args,
            resume_from_checkpoint=resume_from_checkpoint,
            trial=trial,
            ignore_keys_for_eval=ignore_keys_for_eval,
        )

Main training entry point.

Args -----= resume_from_checkpoint (str or bool, optional): If a str, local path to a saved checkpoint as saved by a previous instance of [Trainer]. If a bool and equals True, load the last checkpoint in args.output_dir as saved by a previous instance of [Trainer]. If present, training will resume from the model/optimizer/scheduler states loaded here. trial (optuna.Trial or dict[str, Any], optional): The trial run or the hyperparameter dictionary for hyperparameter search. ignore_keys_for_eval (list[str], optional) A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions for evaluation during the training.

Returns -----= [~trainer_utils.TrainOutput]: Object containing the global step count, training loss, and metrics.

def training_step(self,
model: torch.nn.modules.module.Module,
inputs: dict[str, torch.Tensor | Any],
num_items_in_batch: torch.Tensor | int | None = None) ‑> torch.Tensor

Expand source code

def training_step(
    self,
    model: nn.Module,
    inputs: dict[str, torch.Tensor | Any],
    num_items_in_batch: torch.Tensor | int | None = None,
) -> torch.Tensor:
    """
    Perform a training step on a batch of inputs.

    Subclass and override to inject custom behavior.

    Args:
        model (`nn.Module`):
            The model to train.
        inputs (`dict[str, torch.Tensor | Any]`):
            The inputs and targets of the model.

            The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
            argument `labels`. Check your model's documentation for all accepted arguments.

    Return:
        `torch.Tensor`: The tensor with training loss on this batch.
    """
    # Prepare buffers for context parallelism

    cp_context, inputs = self._prepare_context_parallel_inputs(model, inputs)

    # Context manager is no-op if CP isn't enabled
    with cp_context():
        model.train()
        if hasattr(self.optimizer, "train") and callable(self.optimizer.train):
            self.optimizer.train()

        inputs = self._prepare_inputs(inputs)
        if is_sagemaker_mp_enabled():
            loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
            return loss_mb.reduce_mean().detach().to(self.args.device)

        with self.compute_loss_context_manager():
            loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)

        del inputs
        if (
            self.args.torch_empty_cache_steps is not None
            and self.state.global_step % self.args.torch_empty_cache_steps == 0
        ):
            clear_device_cache()

        kwargs = {}

        # For LOMO optimizers you need to explicitly use the learning rate
        if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
            kwargs["learning_rate"] = self._get_learning_rate()

        if self.args.n_gpu > 1:
            loss = loss.mean()  # mean() to average on multi-gpu parallel training

        # Finally we need to normalize the loss for reporting if GA loss bug is not fixed during compute loss
        if (not self.model_accepts_loss_kwargs or num_items_in_batch is None) and self.compute_loss_func is None:
            # If the model does not accept loss kwargs, we need to normalize the loss by the number of gradient accumulation steps
            loss = loss / self.current_gradient_accumulation_steps

        # Turning off loss scaling w.r.t. gradient accumulation when DeepSpeed is enabled
        # https://github.com/huggingface/transformers/pull/35808
        if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
            kwargs["scale_wrt_gas"] = False

        self.accelerator.backward(loss, **kwargs)

        return loss.detach()

Perform a training step on a batch of inputs.

Subclass and override to inject custom behavior.

Args -----= model (nn.Module): The model to train. inputs (dict[str, torch.Tensor | Any]): The inputs and targets of the model.

The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
argument <code>labels</code>. Check your model's documentation for all accepted arguments.

Return -----= torch.Tensor: The tensor with training loss on this batch.

class TrainingArguments (output_dir: str | None = None,
per_device_train_batch_size: int = 8,
num_train_epochs: float = 3.0,
max_steps: int = -1,
learning_rate: float = 5e-05,
lr_scheduler_type: transformers.trainer_utils.SchedulerType | str = 'linear',
lr_scheduler_kwargs: dict | str | None = None,
warmup_steps: float = 0,
optim: transformers.training_args.OptimizerNames | str = 'adamw_torch_fused',
optim_args: str | None = None,
weight_decay: float = 0.0,
adam_beta1: float = 0.9,
adam_beta2: float = 0.999,
adam_epsilon: float = 1e-08,
optim_target_modules: str | list[str] | None = None,
gradient_accumulation_steps: int = 1,
average_tokens_across_devices: bool = True,
max_grad_norm: float = 1.0,
label_smoothing_factor: float = 0.0,
bf16: bool = False,
fp16: bool = False,
bf16_full_eval: bool = False,
fp16_full_eval: bool = False,
tf32: bool | None = None,
gradient_checkpointing: bool = False,
gradient_checkpointing_kwargs: dict[str, typing.Any] | str | None = None,
torch_compile: bool = False,
torch_compile_backend: str | None = None,
torch_compile_mode: str | None = None,
use_liger_kernel: bool = False,
liger_kernel_config: dict[str, bool] | None = None,
use_cache: bool = False,
neftune_noise_alpha: float | None = None,
torch_empty_cache_steps: int | None = None,
auto_find_batch_size: bool = False,
logging_strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps',
logging_steps: float = 500,
logging_first_step: bool = False,
log_on_each_node: bool = True,
logging_nan_inf_filter: bool = True,
include_num_input_tokens_seen: str | bool = 'no',
log_level: str = 'passive',
log_level_replica: str = 'warning',
disable_tqdm: bool | None = None,
report_to: str | list[str] | None = 'none',
run_name: str | None = None,
project: str = 'huggingface',
trackio_space_id: str | None = None,
trackio_bucket_id: str | None = None,
trackio_static_space_id: str | Literal[False] | None = None,
eval_strategy: transformers.trainer_utils.IntervalStrategy | str = 'no',
eval_steps: float | None = None,
eval_delay: float = 0,
per_device_eval_batch_size: int = 8,
prediction_loss_only: bool = False,
eval_on_start: bool = False,
eval_do_concat_batches: bool = True,
eval_use_gather_object: bool = False,
eval_accumulation_steps: int | None = None,
include_for_metrics: list[str] = <factory>,
batch_eval_metrics: bool = False,
save_only_model: bool = False,
save_strategy: transformers.trainer_utils.SaveStrategy | str = 'steps',
save_steps: float = 500,
save_on_each_node: bool = False,
save_total_limit: int | None = None,
enable_jit_checkpoint: bool = False,
push_to_hub: bool = False,
hub_token: str | None = None,
hub_private_repo: bool | None = None,
hub_model_id: str | None = None,
hub_strategy: transformers.trainer_utils.HubStrategy | str = 'every_save',
hub_always_push: bool = False,
hub_revision: str | None = None,
load_best_model_at_end: bool = False,
metric_for_best_model: str | None = None,
greater_is_better: bool | None = None,
ignore_data_skip: bool = False,
restore_callback_states_from_checkpoint: bool = False,
full_determinism: bool = False,
seed: int = 42,
data_seed: int | None = None,
use_cpu: bool = False,
accelerator_config: dict | str | None = None,
parallelism_config: accelerate.parallelism_config.ParallelismConfig | None = None,
dataloader_drop_last: bool = False,
dataloader_num_workers: int = 0,
dataloader_pin_memory: bool = True,
dataloader_persistent_workers: bool = False,
dataloader_prefetch_factor: int | None = None,
remove_unused_columns: bool = True,
label_names: list[str] | None = None,
train_sampling_strategy: str = 'random',
length_column_name: str = 'length',
ddp_find_unused_parameters: bool | None = None,
ddp_bucket_cap_mb: int | None = None,
ddp_broadcast_buffers: bool | None = None,
ddp_static_graph: bool | None = None,
ddp_backend: str | None = None,
ddp_timeout: int = 1800,
fsdp: list[transformers.trainer_utils.FSDPOption] | str | None = None,
fsdp_config: dict[str, typing.Any] | str | None = None,
deepspeed: dict | str | None = None,
debug: str | list[transformers.debug_utils.DebugOption] = '',
skip_memory_metrics: bool = True,
do_train: bool = False,
do_eval: bool = False,
do_predict: bool = False,
resume_from_checkpoint: str | None = None,
warmup_ratio: float | None = None,
logging_dir: str | None = None,
local_rank: int = -1)

Expand source code

@dataclass
class TrainingArguments:
    """
    Configuration class for controlling all aspects of model training with the Trainer.
    TrainingArguments centralizes all hyperparameters, optimization settings, logging preferences, and infrastructure choices needed for training.

    [`HfArgumentParser`] can turn this class into
    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the
    command line.

    Parameters:
        output_dir (`str`, *optional*, defaults to `"trainer_output"`):
            The output directory where the model predictions and checkpoints will be written.

        > Training Duration and Batch Size

        per_device_train_batch_size (`int`, *optional*, defaults to 8):
            The batch size *per device*. The **global batch size** is computed as:
            `per_device_train_batch_size * number_of_devices` in multi-GPU or distributed setups.
        num_train_epochs(`float`, *optional*, defaults to 3.0):
            Total number of training epochs to perform (if not an integer, will perform the decimal part percents of
            the last epoch before stopping training).
        max_steps (`int`, *optional*, defaults to -1):
            Overrides `num_train_epochs`. If set to a positive number, the total number of training steps to perform.
            For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until
            `max_steps` is reached.

        > Learning Rate & Scheduler

        learning_rate (`float`, *optional*, defaults to 5e-5):
            The initial learning rate for the optimizer. This is typically the peak learning rate when using a scheduler with warmup.
        lr_scheduler_type (`str` or [`SchedulerType`], *optional*, defaults to `"linear"`):
            The learning rate scheduler type to use. See [`SchedulerType`] for all possible values. Common choices:
                - "linear" = [`get_linear_schedule_with_warmup`]
                - "cosine" = [`get_cosine_schedule_with_warmup`]
                - "constant" =  [`get_constant_schedule`]
                - "constant_with_warmup" = [`get_constant_schedule_with_warmup`]
        lr_scheduler_kwargs (`dict` or `str`, *optional*, defaults to `None`):
            The extra arguments for the lr_scheduler. See the documentation of each scheduler for possible values.
        warmup_steps (`int` or `float`, *optional*, defaults to 0):
            Number of steps for a linear warmup from 0 to `learning_rate`. Warmup helps stabilize training in the initial phase. Can be:
                - An integer: exact number of warmup steps
                - A float in range [0, 1): interpreted as ratio of total training steps

        > Optimizer

        optim (`str` or [`training_args.OptimizerNames`], *optional*, defaults to `"adamw_torch"` (for torch>=2.8 `"adamw_torch_fused"`)):
            The optimizer to use. Common options:
                - `"adamw_torch"`: PyTorch's AdamW (recommended default)
                - `"adamw_torch_fused"`: Fused AdamW kernel
                - `"adamw_hf"`: HuggingFace's AdamW implementation
                - `"sgd"`: Stochastic Gradient Descent with momentum
                - `"adafactor"`: Memory-efficient optimizer for large models
                - `"adamw_8bit"`: 8-bit AdamW (requires bitsandbytes)
                See [`OptimizerNames`] for the complete list.
        optim_args (`str`, *optional*):
            Optional arguments that are supplied to optimizers such as AnyPrecisionAdamW, AdEMAMix, and GaLore.
        weight_decay (`float`, *optional*, defaults to 0):
            Weight decay coefficient applied by the optimizer (not the loss function). Adds L2
            regularization to prevent overfitting by penalizing large weights. Automatically
            excluded from bias and LayerNorm parameters. Typical values: 0.01 (standard), 0.1
            (stronger regularization), 0.0 (no regularization).
        adam_beta1 (`float`, *optional*, defaults to 0.9):
            The exponential decay rate for the first moment estimates (momentum) in Adam-based
            optimizers. Controls how much history of gradients to retain.
        adam_beta2 (`float`, *optional*, defaults to 0.999):
            The exponential decay rate for the second moment estimates (variance) in Adam-based
            optimizers. Controls adaptive learning rate scaling.
        adam_epsilon (`float`, *optional*, defaults to 1e-8):
            Epsilon value for numerical stability in Adam-based optimizers. Prevents division by
            zero in the denominator of the update rule.
        optim_target_modules (`Union[str, list[str]]`, *optional*):
            The target modules to optimize, i.e. the module names that you would like to train.
            Currently used for the [GaLore algorithm](https://huggingface.co/papers/2403.03507) and [APOLLO algorithm](https://huggingface.co/papers/2412.05270).
            See [GaLore implementation](https://github.com/jiaweizzhao/GaLore) and [APOLLO implementation](https://github.com/zhuhanqing/APOLLO) for more details.
            You need to make sure to pass a valid GaLore or APOLLO optimizer, e.g., one of: "apollo_adamw", "galore_adamw", "galore_adamw_8bit", "galore_adafactor" and make sure that the target modules are `nn.Linear` modules only.

        > Regularization & Training Stability

        gradient_accumulation_steps (`int`, *optional*, defaults to 1):
            Number of update steps to accumulate gradients before performing a backward/update pass.
            Simulates larger batch sizes without additional memory. Effective batch size =
            `per_device_train_batch_size × num_devices × gradient_accumulation_steps`.
            > [!TIP]
            > When using gradient accumulation, one "step" is counted as one step with a backward pass. Therefore, logging, evaluation, and saving will occur every `gradient_accumulation_steps × xxx_step` training examples.
        average_tokens_across_devices (`bool`, *optional*, defaults to `True`):
            Whether or not to average tokens across devices. If enabled, will use all_reduce to synchronize
            num_tokens_in_batch for precise loss calculation. Reference:
            https://github.com/huggingface/transformers/issues/34242
        max_grad_norm (`float`, *optional*, defaults to 1.0):
            Maximum gradient norm for gradient clipping. Applied after backward pass, before
            optimizer step. Prevents gradient explosion by scaling down gradients when their global
            norm exceeds this threshold. Set to 0 to disable clipping. Typical values:
            1.0 (standard), 0.5 (more conservative), 5.0 (less aggressive).
        label_smoothing_factor (`float`, *optional*, defaults to 0.0):
            Label smoothing factor to prevent overconfidence. Replaces hard 0/1 targets with soft
            targets: 0 becomes `ε/num_labels` and 1 becomes `1 - ε + ε/num_labels`, where
            ε = `label_smoothing_factor`. Zero means no smoothing. Typical range: 0.0 to 0.1.

        > Mixed Precision Training

        bf16 (`bool`, *optional*, defaults to `False`):
            Enable bfloat16 (BF16) mixed precision training
            Generally preferred over FP16 due to better numerical stability and no loss scaling required.
        fp16 (`bool`, *optional*, defaults to `False`):
            Enable float16 (FP16) mixed precision training.
            Consider using BF16 instead if your hardware supports it.
        bf16_full_eval (`bool`, *optional*, defaults to `False`):
            Use full BF16 precision for evaluation (not just mixed precision). Faster and saves
            memory but may affect metric values slightly. Only applies during evaluation.
        fp16_full_eval (`bool`, *optional*, defaults to `False`):
            Use full FP16 precision for evaluation (not just mixed precision). Faster and saves
            memory but may affect metric values slightly. Only applies during evaluation.
        tf32 (`bool`, *optional*):
            Enable TensorFloat-32 (TF32) mode on Ampere and newer GPUs. TF32 uses 19-bit precision
            for matrix multiplications (instead of FP32's 23-bit), providing up to 8x speedup with
            negligible accuracy loss. Default depends on PyTorch version. See
            [TF32 docs](https://huggingface.co/docs/transformers/perf_train_gpu_one#tf32).

        > Gradient Checkpointing

        gradient_checkpointing (`bool`, *optional*, defaults to `False`):
            Enable gradient checkpointing to trade compute for memory. Reduces memory usage by
            clearing activations during forward pass and recomputing them during backward pass.
            Enables training larger models or batch sizes at the cost of ~20% slower training.
        gradient_checkpointing_kwargs (`dict`, *optional*, defaults to `None`):
            Keyword arguments passed to `gradient_checkpointing_enable()`.

        > Compilation

        torch_compile (`bool`, *optional*, defaults to `False`):
            Compile the model using PyTorch 2.0's `torch.compile()` for faster training. Can provide
            20-50% speedup with no code changes. Uses default compilation settings unless
            `torch_compile_backend` or `torch_compile_mode` are specified.
        torch_compile_backend (`str`, *optional*):
            Backend for `torch.compile()`. If set, automatically enables `torch_compile`. Options
            include `"inductor"` (default), `"aot_eager"`, `"cudagraphs"`. Backends vary by PyTorch
            version - see PyTorch docs for available options.
        torch_compile_mode (`str`, *optional*):
            Compilation mode for `torch.compile()`. If set, automatically enables `torch_compile`.
            Options: `"default"`, `"reduce-overhead"` (minimize Python overhead), `"max-autotune"`
            (aggressive optimization, slower compile time).

        > Kernels

        use_liger_kernel (`bool`, *optional*, defaults to `False`):
            Enable [Liger Kernel](https://github.com/linkedin/Liger-Kernel) optimizations. Increases
            multi-GPU throughput by ~20% and reduces memory usage by ~60%. Works with Flash Attention,
            FSDP, and DeepSpeed. Currently supports Llama, Mistral, Mixtral, and Gemma models.
        liger_kernel_config (`Optional[dict]`, *optional*):
            Configuration for Liger Kernel. Passed as kwargs to `_apply_liger_kernel_to_instance()`.
            Options typically include: `"rope"`, `"swiglu"`, `"cross_entropy"`,
            `"fused_linear_cross_entropy"`, `"rms_norm"`. If `None`, uses default configuration.

        > Additional Optimizations

        use_cache (`bool`, *optional*, defaults to `False`):
            Whether or not to enable cache for the model. For training, this is usually not needed apart from some PEFT methods that uses `past_key_values`.
        neftune_noise_alpha (`Optional[float]`):
            If not `None`, this will activate NEFTune noise embeddings. This can drastically improve model performance
            for instruction fine-tuning. Check out the [original paper](https://huggingface.co/papers/2310.05914) and the
            [original code](https://github.com/neelsjain/NEFTune). Support transformers `PreTrainedModel` and also
            `PeftModel` from peft. The original paper used values in the range [5.0, 15.0].
        torch_empty_cache_steps (`int`, *optional*):
            Number of steps to wait before calling `torch.<device>.empty_cache()`. If left unset or set to None, cache will not be emptied.
            This can help avoid CUDA out-of-memory errors by lowering peak VRAM usage at a cost of about [10% slower performance](https://github.com/huggingface/transformers/issues/31372).
        auto_find_batch_size (`bool`, *optional*, defaults to `False`)
            Whether to find a batch size that will fit into memory automatically through exponential decay, avoiding
            CUDA Out-of-Memory errors.

        > Logging & Monitoring Training

        logging_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`):
            The logging strategy to adopt during training. Possible values are:
                - `"no"`: No logging is done during training.
                - `"epoch"`: Logging is done at the end of each epoch.
                - `"steps"`: Logging is done every `logging_steps`.
        logging_steps (`int` or `float`, *optional*, defaults to 500):
            Number of update steps between two logs if `logging_strategy="steps"`. Should be an integer or a float in
            range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps.
        logging_first_step (`bool`, *optional*, defaults to `False`):
            Whether to log the first `global_step` or not.
        log_on_each_node (`bool`, *optional*, defaults to `True`):
            In multinode distributed training, whether to log using `log_level` once per node, or only on the main
            node.
        logging_nan_inf_filter (`bool`, *optional*, defaults to `True`):
             Filter out NaN and Inf losses when logging. If `True`, replaces NaN/Inf losses with the
            average of recent valid losses. Does not affect gradient computation, only logging.
        include_num_input_tokens_seen (`Optional[Union[str, bool]]`, *optional*, defaults to "no"):
            Whether to track the number of input tokens seen. Must be one of ["all", "non_padding", "no"] or a boolean value which map to "all" or "no".
            May be slower in distributed training as gather operations must be called.

        > Logging

        log_level (`str`, *optional*, defaults to `passive`):
            Logging level for the main process. Options: `"debug"`, `"info"`, `"warning"`, `"error"`,
            `"critical"`, or `"passive"` (doesn't change the current Transformers logging level,
            which defaults to `"warning"`)
        log_level_replica (`str`, *optional*, defaults to `"warning"`):
            Logging level for replica processes in distributed training. Same options as `log_level`.
        disable_tqdm (`bool`, *optional*):
            Disable tqdm progress bars. Defaults to `True` if `log_level` is warning or lower, `False` otherwise.

        > Experiment Tracking Integration

        report_to (`str` or `list[str]`, *optional*, defaults to `"none"`):
            The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`,
            `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`, `"swanlab"`,
            `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations installed, `"none"`
            for no integrations.
        run_name (`str`, *optional*):
            A descriptor for the run. Typically used for [trackio](https://github.com/gradio-app/trackio),
            [wandb](https://www.wandb.com/), [mlflow](https://www.mlflow.org/), [comet](https://www.comet.com/site) and
            [swanlab](https://swanlab.cn) logging.
        project (`str`, *optional*, defaults to `"huggingface"`):
            The name of the project to use for logging. Currently, only used by Trackio.
        trackio_space_id (`str` or `None`, *optional*, defaults to `None`):
            The Hugging Face Space ID to use for live Trackio logging with a Gradio-based Space. Should be a full
            Space name like `'username/reponame'` or `'orgname/reponame'`, or just `'reponame'` (the Space is created in
            the currently logged-in user's namespace). If `None`, metrics are logged only to a **local** directory (no
            Space on the Hub). That Gradio Space has **read and write** access to the Trackio **Bucket**, which is what you want while
            training is **in progress**—for example when **resuming** a partial run or **aggregating logs** from multiple
            machines. The Space will be **public** unless you set `hub_private_repo=True` or your organization's default is to
            create private Spaces.
        trackio_bucket_id (`str` or `None`, *optional*, defaults to `None`):
            Optional Hugging Face Bucket id for Trackio. If unset, Trackio derives one. Used together with a Gradio Space
            (`trackio_space_id`) and when deploying a static Space (`trackio_static_space_id` is not `False`).
        trackio_static_space_id (`str`, `False`, or `None`, *optional*, defaults to `None`):
            The Hugging Face Space ID to use for static Space created after training is complete. Should be a full
            Space name like `'username/reponame'` or `'orgname/reponame'`, or just `'reponame'` (the Space is created in
            the currently logged-in user's namespace). If False, no static Space will be created. If None, and model is pushed to the Hub,
            a static Space will be created with a default name and this will be linked from the model card. The Space will be public
            unless you set `hub_private_repo=True` or your organization's default is to create private Spaces.

        > Evaluation

        eval_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"no"`):
            When to run evaluation. Options:
            - `"no"`: No evaluation during training
            - `"steps"`: Evaluate every `eval_steps`
            - `"epoch"`: Evaluate at the end of each epoch
        eval_steps (`int` or `float`, *optional*):
            Number of update steps between two evaluations if `eval_strategy="steps"`. Will default to the same
            value as `logging_steps` if not set. Should be an integer or a float in range `[0,1)`. If smaller than 1,
            will be interpreted as ratio of total training steps.
        eval_delay (`float`, *optional*):
            Number of epochs or steps to wait for before the first evaluation can be performed, depending on the
            eval_strategy.
        per_device_eval_batch_size (`int`, *optional*, defaults to 8):
            The batch size per device accelerator core/CPU for evaluation.
        prediction_loss_only (`bool`, *optional*, defaults to `False`):
            When performing evaluation and generating predictions, only returns the loss.
        eval_on_start (`bool`, *optional*, defaults to `False`):
            Whether to perform a evaluation step (sanity check) before the training to ensure the validation steps works correctly.
        eval_do_concat_batches (`bool`, *optional*, defaults to `True`):
            Whether to recursively concat inputs/losses/labels/predictions across batches. If `False`,
            will instead store them as lists, with each batch kept separate.
        eval_use_gather_object (`bool`, *optional*, defaults to `False`):
            Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices. This should only be enabled if users are not just returning tensors, and this is actively discouraged by PyTorch.
            This is useful when the labels structure is non standard, like in computer vision tasks.
        eval_accumulation_steps (`int`, *optional*):
            Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. If
            left unset, the whole predictions are accumulated on the device accelerator before being moved to the CPU (faster but
            requires more memory).

        > Metrics Computation

        include_for_metrics (`list[str]`, *optional*, defaults to `[]`):
            Include additional data in the `compute_metrics` function if needed for metrics computation.
            Possible options to add to `include_for_metrics` list:
            - `"inputs"`: Input data passed to the model, intended for calculating input dependent metrics.
            - `"loss"`: Loss values computed during evaluation, intended for calculating loss dependent metrics.
        batch_eval_metrics (`bool`, *optional*, defaults to `False`):
            If set to `True`, evaluation will call compute_metrics at the end of each batch to accumulate statistics
            rather than saving all eval logits in memory. When set to `True`, you must pass a compute_metrics function
            that takes a boolean argument `compute_result`, which when passed `True`, will trigger the final global
            summary statistics from the batch-level summary statistics you've accumulated over the evaluation set.

        > Checkpointing & Saving

        save_only_model (`bool`, *optional*, defaults to `False`):
            Save only model weights, not optimizer/scheduler/RNG state. Significantly reduces
            checkpoint size but prevents resuming training from the checkpoint. Use when you only
            need the trained model for inference, not continued training.
            You can only load the model using `from_pretrained` with this option set to `True`.
        save_strategy (`str` or [`~trainer_utils.SaveStrategy`], *optional*, defaults to `"steps"`):
            The checkpoint save strategy to adopt during training. Possible values are:
                - `"no"`: No save is done during training.
                - `"epoch"`: Save is done at the end of each epoch.
                - `"steps"`: Save is done every `save_steps`.
                - `"best"`: Save is done whenever a new `best_metric` is achieved.
        save_steps (`int` or `float`, *optional*, defaults to 500):
            Number of updates steps before two checkpoint saves if `save_strategy="steps"`. Should be an integer or a
            float in range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps.
        save_on_each_node (`bool`, *optional*, defaults to `False`):
            When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on
            the main one.
            This should not be activated when the different nodes use the same storage as the files will be saved with
            the same names for each node.
        save_total_limit (`int`, *optional*):
            Maximum number of checkpoints to keep. Deletes older checkpoints in `output_dir`. When
            `load_best_model_at_end=True`, the best checkpoint is always retained plus the most
            recent ones. For example, `save_total_limit=5` keeps the 4 most recent plus the best
        enable_jit_checkpoint (`bool`, *optional*, defaults to `False`):
            Enable Just-In-Time checkpointing on SIGTERM signal for graceful termination on
            preemptible workloads. **Important**: Configure your orchestrator's graceful shutdown
            period to allow sufficient time. For Kubernetes, set `terminationGracePeriodSeconds`
            (default 30s is usually insufficient). For Slurm, use `--signal=USR1@<seconds>`.
            Required grace period ≥ longest iteration time + checkpoint save time.

        > Hugging Face Hub Integration

        push_to_hub (`bool`, *optional*, defaults to `False`):
            Whether or not to push the model to the Hub every time the model is saved. If this is activated,
            `output_dir` will begin a git directory synced with the repo (determined by `hub_model_id`) and the content
            will be pushed each time a save is triggered (depending on your `save_strategy`). Calling
            [`~Trainer.save_model`] will also trigger a push.
        hub_token (`str`, *optional*):
            The token to use to push the model to the Hub. Will default to the token in the cache folder obtained with
            `hf auth login`.
        hub_private_repo (`bool`, *optional*):
            Whether to make the repo private. If `None` (default), the repo will be public unless the organization's
            default is private. This value is ignored if the repo already exists. If reporting to Trackio with
            deployment to Hugging Face Spaces enabled, the same logic determines whether the Space is private.
        hub_model_id (`str`, *optional*):
            The name of the repository to keep in sync with the local *output_dir*. It can be a simple model ID in
            which case the model will be pushed in your namespace. Otherwise it should be the whole repository name,
            for instance `"user_name/model"`, which allows you to push to an organization you are a member of with
            `"organization_name/model"`. Will default to `user_name/output_dir_name` with *output_dir_name* being the
            name of `output_dir`.
        hub_strategy (`str` or [`~trainer_utils.HubStrategy`], *optional*, defaults to `"every_save"`):
            Defines what and when to push to Hub. Options:
                - `"end"`: Push only at the end of training
                - `"every_save"`: Push on each save (async to not block training)
                - `"checkpoint"`: Like `"every_save"` plus push latest checkpoint to `"last-checkpoint"` subfolder for easy resuming
                - `"all_checkpoints"`: Push all checkpoints as they appear
        hub_always_push (`bool`, *optional*, defaults to `False`):
            Unless this is `True`, the `Trainer` will skip pushing a checkpoint when the previous push is not finished.
        hub_revision (`str`, *optional*):
            The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash.

        > Best Model Tracking

        load_best_model_at_end (`bool`, *optional*, defaults to `False`):
            Load the best checkpoint at the end of training. Requires `eval_strategy` to be set.
            When enabled, the best checkpoint is always saved (see `save_total_limit`).
            <Tip>
            When `True`, `save_strategy` must match `eval_strategy` (unless `save_strategy` is `"best"`), and if using `"steps"`,
            `save_steps` must be a multiple of `eval_steps`.
            </Tip>
        metric_for_best_model (`str`, *optional*):
            Metric to use for comparing models when `load_best_model_at_end=True`. Must be a metric
            name returned by evaluation, with or without the `"eval_"` prefix. Defaults to `"loss"`.
            If you set this, `greater_is_better` will default to `True` unless the name ends with
            `"loss"`. Examples: `"accuracy"`, `"f1"`, `"eval_bleu"`.
        greater_is_better (`bool`, *optional*):
            Whether higher metric values are better. Defaults based on `metric_for_best_model`:
            `True` if the metric name doesn't end in `"loss"`, `False` otherwise.

        > Resuming Training

        ignore_data_skip (`bool`, *optional*, defaults to `False`):
            When resuming training, skip fast-forwarding through the dataset to reach the previous
            state. If `True`, training starts from the beginning of the dataset (faster resume but
            results won't match interrupted training). If `False`, skips seen data (slower resume
            but exact continuation).
        restore_callback_states_from_checkpoint (`bool`, *optional*, defaults to `False`):
            Restore callback states from checkpoint when resuming. If `True`, will override callbacks
            passed to Trainer if they exist in the checkpoint.

        > Reproducibility

        full_determinism (`bool`, *optional*, defaults to `False`)
            If `True`, [`enable_full_determinism`] is called instead of [`set_seed`] to ensure reproducible results in
            distributed training. Important: this will negatively impact the performance, so only use it for debugging.
        seed (`int`, *optional*, defaults to 42):
            Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use the
            [`~Trainer.model_init`] function to instantiate the model if it has some randomly initialized parameters.
        data_seed (`int`, *optional*):
            Random seed to be used with data samplers. If not set, random generators for data sampling will use the
            same seed as `seed`. This can be used to ensure reproducibility of data sampling, independent of the model
            seed.

        > Hardware Configuration

        use_cpu (`bool`, *optional*, defaults to `False`):
            Whether or not to use cpu. If set to False, we will use the available torch device/backend.

        > Accelerate Configuration

        accelerator_config (`str`, `dict`, or `AcceleratorConfig`, *optional*):
            Configuration for the internal Accelerate integration. Can be:
                - Path to JSON config file: `"accelerator_config.json"`
                - Dictionary with config options
                - `AcceleratorConfig` instance
            Key options:
                - `split_batches` (`bool`, defaults to `False`): Whether to split batches across devices.
                    If `True`, actual batch size is the same on all devices (total must be divisible by
                    num_processes). If `False`, each device gets the specified batch size.
                - `dispatch_batches` (`bool`): If `True`, only main process iterates through dataloader
                    and dispatches batches to devices. Defaults to `True` for `IterableDataset`, `False`
                    otherwise.
                - `even_batches` (`bool`, defaults to `True`): Duplicate samples from dataset start to
                    ensure all workers get equal batch sizes.
                - `use_seedable_sampler` (`bool`, defaults to `True`): Use fully seedable random sampler
                    for reproducibility.
                - `use_configured_state` (`bool`, defaults to `False`): Use pre-initialized
                    `AcceleratorState`/`PartialState` instead of creating new one. May cause issues with
                    hyperparameter tuning.

        parallelism_config (`ParallelismConfig`, *optional*):
            Parallelism configuration for the training run. Requires Accelerate `1.10.1`

        > Dataloader

        dataloader_drop_last (`bool`, *optional*, defaults to `False`):
            Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size)
            or not.
        dataloader_num_workers (`int`, *optional*, defaults to 0):
            Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in the
            main process.
        dataloader_pin_memory (`bool`, *optional*, defaults to `True`):
            Whether you want to pin memory in data loaders or not. Will default to `True`.
        dataloader_persistent_workers (`bool`, *optional*, defaults to `False`):
            If True, the data loader will not shut down the worker processes after a dataset has been consumed once.
            This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will
            increase RAM usage. Will default to `False`.
        dataloader_prefetch_factor (`int`, *optional*):
            Number of batches loaded in advance by each worker.
            2 means there will be a total of 2 * num_workers batches prefetched across all workers.
        remove_unused_columns (`bool`, *optional*, defaults to `True`):
            Whether or not to automatically remove the columns unused by the model forward method.
        label_names (`list[str]`, *optional*):
            The list of keys in your dictionary of inputs that correspond to the labels.
            Will eventually default to the list of argument names accepted by the model that contain the word "label",
            except if the model used is one of the `XxxForQuestionAnswering` in which case it will also include the
            `["start_positions", "end_positions"]` keys.
            You should only specify `label_names` if you're using custom label names or if your model's `forward` consumes multiple label tensors (e.g., extractive QA).
        train_sampling_strategy (`str`, *optional*, defaults to `"random"`):
            The sampler to use for the training dataloader. Possible values are:

                - `"random"`: Uses `RandomSampler` (default).
                - `"sequential"`: Uses `SequentialSampler`.
                - `"group_by_length"`: Uses `LengthGroupedSampler` to group samples of roughly the same length
                  together (to minimize padding and be more efficient).

            Note: When using an `IterableDataset`, this argument is ignored.
        length_column_name (`str`, *optional*, defaults to `"length"`):
            Column name for precomputed lengths. If the column exists, grouping by length will use these values rather
            than computing them on train startup. Ignored unless `train_sampling_strategy` is `"group_by_length"` and the dataset
            is an instance of `Dataset`.

        > DDP (DistributedDataParallel)

        ddp_find_unused_parameters (`bool`, *optional*):
            When using distributed training, the value of the flag `find_unused_parameters` passed to
            `DistributedDataParallel`. Will default to `False` if gradient checkpointing is used, `True` otherwise.
        ddp_bucket_cap_mb (`int`, *optional*):
            When using distributed training, the value of the flag `bucket_cap_mb` passed to `DistributedDataParallel`.
        ddp_broadcast_buffers (`bool`, *optional*):
            When using distributed training, the value of the flag `broadcast_buffers` passed to
            `DistributedDataParallel`. Will default to `False` if gradient checkpointing is used, `True` otherwise.
        ddp_static_graph (`bool`, *optional*):
            When using distributed training, the value of the flag `static_graph` passed to
            `DistributedDataParallel`.
        ddp_backend (`str`, *optional*):
            The backend to use for distributed training. Must be one of `"nccl"`, `"mpi"`, `"xccl"`, `"gloo"`, `"hccl"`.
        ddp_timeout (`int`, *optional*, defaults to 1800):
            The timeout for `torch.distributed.init_process_group` calls, used to avoid GPU socket timeouts when
            performing slow operations in distributed runnings. Please refer to the [PyTorch documentation](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more
            information.

        > FSDP (Fully Sharded Data Parallel)

        fsdp (`bool`, `str` or list of [`~trainer_utils.FSDPOption`], *optional*, defaults to `None`):
            Enable PyTorch Fully Sharded Data Parallel (FSDP) for distributed training. Options:
                - `"full_shard"`: Shard parameters, gradients, and optimizer states (most memory efficient)
                - `"shard_grad_op"`: Shard only optimizer states and gradients (ZeRO-2)
                - `"hybrid_shard"`: Full shard within nodes, replicate across nodes
                - `"hybrid_shard_zero2"`: Shard gradients/optimizer within nodes, replicate across nodes
                - `"offload"`: Offload parameters and gradients to CPU (only with `"full_shard"` or
                    `"shard_grad_op"`)
                - `"auto_wrap"`: Automatically wrap layers using `default_auto_wrap_policy`
        fsdp_config (`str` or `dict`, *optional*):
            Config to be used with fsdp (Pytorch Distributed Parallel Training). The value is either a location of
            fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`.

            A List of config and its options:
                - fsdp_version (`int`, *optional*, defaults to `1`):
                    The version of FSDP to use. Defaults to 1.
                - min_num_params (`int`, *optional*, defaults to `0`):
                    FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `fsdp` field is
                    passed).
                - transformer_layer_cls_to_wrap (`list[str]`, *optional*):
                    List of transformer layer class names (case-sensitive) to wrap, e.g, `BertLayer`, `GPTJBlock`,
                    `T5Block` .... (useful only when `fsdp` flag is passed).
                - backward_prefetch (`str`, *optional*)
                    FSDP's backward prefetch mode. Controls when to prefetch next set of parameters (useful only when
                    `fsdp` field is passed).

                    A list of options along the following:

                    - `"backward_pre"` : Prefetches the next set of parameters before the current set of parameter's
                      gradient computation.
                    - `"backward_post"` : This prefetches the next set of parameters after the current set of
                      parameter's gradient computation.
                - forward_prefetch (`bool`, *optional*, defaults to `False`)
                    FSDP's forward prefetch mode (useful only when `fsdp` field is passed).
                     If `"True"`, then FSDP explicitly prefetches the next upcoming all-gather while executing in the
                     forward pass.
                - limit_all_gathers (`bool`, *optional*, defaults to `False`)
                    FSDP's limit_all_gathers (useful only when `fsdp` field is passed).
                     If `"True"`, FSDP explicitly synchronizes the CPU thread to prevent too many in-flight
                     all-gathers.
                - use_orig_params (`bool`, *optional*, defaults to `True`)
                    If `"True"`, allows non-uniform `requires_grad` during init, which means support for interspersed
                    frozen and trainable parameters. Useful in cases such as parameter-efficient fine-tuning. Please
                    refer this
                    [blog](https://dev-discuss.pytorch.org/t/rethinking-pytorch-fully-sharded-data-parallel-fsdp-from-first-principles/1019
                - sync_module_states (`bool`, *optional*, defaults to `True`)
                    If `"True"`, each individually wrapped FSDP unit will broadcast module parameters from rank 0 to
                    ensure they are the same across all ranks after initialization
                - cpu_ram_efficient_loading (`bool`, *optional*, defaults to `False`)
                    If `"True"`, only the first process loads the pretrained model checkpoint while all other processes
                    have empty weights.  When this setting as `"True"`, `sync_module_states` also must to be `"True"`,
                    otherwise all the processes except the main process would have random weights leading to unexpected
                    behaviour during training.
                - activation_checkpointing (`bool`, *optional*, defaults to `False`):
                    If `"True"`, activation checkpointing is a technique to reduce memory usage by clearing activations of
                    certain layers and recomputing them during a backward pass. Effectively, this trades extra
                    computation time for reduced memory usage.
                - xla (`bool`, *optional*, defaults to `False`):
                    Whether to use PyTorch/XLA Fully Sharded Data Parallel Training. This is an experimental feature
                    and its API may evolve in the future.
                - xla_fsdp_settings (`dict`, *optional*)
                    The value is a dictionary which stores the XLA FSDP wrapping parameters.

                    For a complete list of options, please see [here](
                    https://github.com/pytorch/xla/blob/master/torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py).
                - xla_fsdp_grad_ckpt (`bool`, *optional*, defaults to `False`):
                    Will use gradient checkpointing over each nested XLA FSDP wrapped layer. This setting can only be
                    used when the xla flag is set to true, and an auto wrapping policy is specified through
                    fsdp_min_num_params or fsdp_transformer_layer_cls_to_wrap.

        > DeepSpeed

        deepspeed (`str` or `dict`, *optional*):
             Enable [DeepSpeed](https://github.com/deepspeedai/DeepSpeed) integration. Value is either:
                - Path to DeepSpeed JSON config file: `"ds_config.json"`
                - Loaded config as dictionary
          > [!TIP]
          > If using ZeRO initialization, instantiate your model *after* initializing
          `TrainingArguments`, otherwise ZeRO won't be applied.

        > Debugging & Profiling (Experimental)

        debug (`str` or list of [`~debug_utils.DebugOption`], *optional*, defaults to `""`):
            Enable one or more debug features. This is an experimental feature.
            Possible options are:
            - "underflow_overflow": detects overflow in model's input/outputs and reports the last frames that led to
              the event
            - "tpu_metrics_debug": print debug metrics on TPU
        skip_memory_metrics (`bool`, *optional*, defaults to `True`):
            Whether to skip adding of memory profiler reports to metrics. This is skipped by default because it slows
            down the training and evaluation speed.

        > External Script Flags (not used by Trainer)

        do_train (`bool`, *optional*, defaults to `False`):
            Whether to run training or not. This argument is not directly used by [`Trainer`], it's intended to be used
            by your training/evaluation scripts instead. See the [example
            scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details.
        do_eval (`bool`, *optional*):
            Whether to run evaluation on the validation set or not. Will be set to `True` if `eval_strategy` is
            different from `"no"`. This argument is not directly used by [`Trainer`], it's intended to be used by your
            training/evaluation scripts instead. See the [example
            scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details.
        do_predict (`bool`, *optional*, defaults to `False`):
            Whether to run predictions on the test set or not. This argument is not directly used by [`Trainer`], it's
            intended to be used by your training/evaluation scripts instead. See the [example
            scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details.
        resume_from_checkpoint (`str`, *optional*):
            The path to a folder with a valid checkpoint for your model. This argument is not directly used by
            [`Trainer`], it's intended to be used by your training/evaluation scripts instead. See the [example
            scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details.
    """

    # Fields that accept dict values via CLI as JSON strings (e.g., '{"key": "value"}').
    # Any new dict-typed arg must be added here and typed as `dict | str | None`.
    _VALID_DICT_FIELDS = [
        "accelerator_config",
        "fsdp_config",
        "deepspeed",
        "gradient_checkpointing_kwargs",
        "lr_scheduler_kwargs",
    ]

    # --- Output ---
    output_dir: str | None = field(
        default=None,
        metadata={"help": "The output directory where the model predictions and checkpoints will be written."},
    )

    # --- Training Duration and Batch Size ---
    per_device_train_batch_size: int = field(default=8, metadata={"help": "The batch size per device for training."})
    num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."})
    max_steps: int = field(
        default=-1,
        metadata={
            "help": "Overrides `num_train_epochs`. If set to a positive number, the total number of training steps to perform."
        },
    )

    # --- Learning Rate & Scheduler ---
    learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for the optimizer."})
    lr_scheduler_type: SchedulerType | str = field(
        default="linear",
        metadata={"help": "The learning rate scheduler type to use. See `SchedulerType` for all possible values."},
    )
    lr_scheduler_kwargs: dict | str | None = field(
        default=None,
        metadata={
            "help": "The extra arguments for the lr_scheduler. See the documentation of each scheduler for possible values."
        },
    )
    warmup_steps: float = field(
        default=0,
        metadata={
            "help": "Number of steps for a linear warmup from 0 to `learning_rate`. Can be an integer (exact steps) or a float in [0, 1) (ratio of total steps)."
        },
    )

    # --- Optimizer ---
    default_optim = "adamw_torch"
    if is_torch_available():
        from .pytorch_utils import is_torch_greater_or_equal_than_2_8

        if is_torch_greater_or_equal_than_2_8:
            default_optim = "adamw_torch_fused"
    optim: OptimizerNames | str = field(
        default=default_optim,
        metadata={"help": "The optimizer to use. See `OptimizerNames` for the complete list."},
    )
    optim_args: str | None = field(
        default=None,
        metadata={
            "help": "Optional arguments supplied to optimizers such as AnyPrecisionAdamW, AdEMAMix, and GaLore."
        },
    )
    weight_decay: float = field(
        default=0.0,
        metadata={
            "help": "Weight decay coefficient applied by the optimizer. Automatically excluded from bias and LayerNorm parameters."
        },
    )
    adam_beta1: float = field(
        default=0.9,
        metadata={
            "help": "The exponential decay rate for the first moment estimates (momentum) in Adam-based optimizers."
        },
    )
    adam_beta2: float = field(
        default=0.999,
        metadata={
            "help": "The exponential decay rate for the second moment estimates (variance) in Adam-based optimizers."
        },
    )
    adam_epsilon: float = field(
        default=1e-8, metadata={"help": "Epsilon value for numerical stability in Adam-based optimizers."}
    )
    optim_target_modules: None | str | list[str] = field(
        default=None,
        metadata={"help": "The target modules to optimize. Currently used for the GaLore and APOLLO algorithms."},
    )

    # --- Regularization & Training Stability ---
    gradient_accumulation_steps: int = field(
        default=1,
        metadata={
            "help": (
                "Number of update steps to accumulate gradients before performing a backward/update pass."
                " Effective batch size = per_device_train_batch_size * num_devices * gradient_accumulation_steps."
            )
        },
    )
    average_tokens_across_devices: bool = field(
        default=True,
        metadata={
            "help": "Whether or not to average tokens across devices. If enabled, will use all_reduce to "
            "synchronize num_tokens_in_batch for precise loss calculation. Reference: "
            "https://github.com/huggingface/transformers/issues/34242"
        },
    )
    max_grad_norm: float = field(
        default=1.0, metadata={"help": "Maximum gradient norm for gradient clipping. Set to 0 to disable."}
    )
    label_smoothing_factor: float = field(
        default=0.0, metadata={"help": "Label smoothing factor to prevent overconfidence. Zero means no smoothing."}
    )

    # --- Mixed Precision ---
    bf16: bool = field(
        default=False,
        metadata={
            "help": "Enable bfloat16 (BF16) mixed precision training. Generally preferred over FP16 due to better numerical stability."
        },
    )
    fp16: bool = field(
        default=False,
        metadata={
            "help": "Enable float16 (FP16) mixed precision training. Consider using BF16 instead if your hardware supports it."
        },
    )
    bf16_full_eval: bool = field(
        default=False,
        metadata={
            "help": "Use full BF16 precision for evaluation (not just mixed precision). Faster and saves memory."
        },
    )
    fp16_full_eval: bool = field(
        default=False,
        metadata={
            "help": "Use full FP16 precision for evaluation (not just mixed precision). Faster and saves memory."
        },
    )
    tf32: bool | None = field(
        default=None,
        metadata={
            "help": "Enable TF32 mode on Ampere and newer GPUs. Provides up to 8x speedup with negligible accuracy loss."
        },
    )

    # --- Gradient Checkpointing ---
    gradient_checkpointing: bool = field(
        default=False,
        metadata={
            "help": "Enable gradient checkpointing to trade compute for memory. Reduces memory at the cost of ~20%% slower training."
        },
    )
    gradient_checkpointing_kwargs: dict[str, Any] | str | None = field(
        default=None,
        metadata={"help": "Keyword arguments passed to `gradient_checkpointing_enable()`."},
    )

    # --- Compilation ---
    torch_compile: bool = field(
        default=False, metadata={"help": "Compile the model using `torch.compile()` for faster training."}
    )
    torch_compile_backend: str | None = field(
        default=None,
        metadata={
            "help": "Backend for `torch.compile()`. If set, automatically enables `torch_compile`.",
        },
    )
    torch_compile_mode: str | None = field(
        default=None,
        metadata={
            "help": "Compilation mode for `torch.compile()`. If set, automatically enables `torch_compile`.",
        },
    )

    # --- Kernels ---
    use_liger_kernel: bool = field(
        default=False,
        metadata={
            "help": "Enable Liger Kernel optimizations. Increases throughput by ~20%% and reduces memory by ~60%%."
        },
    )
    liger_kernel_config: dict[str, bool] | None = field(
        default=None,
        metadata={
            "help": "Configuration for Liger Kernel. Passed as kwargs to `_apply_liger_kernel_to_instance()`. If None, uses default configuration."
        },
    )

    # --- Additional Optimizations ---
    use_cache: bool = field(
        default=False,
        metadata={
            "help": "Whether or not to use cache for the model For training, this is usually not needed apart from some PEFT methods that uses `past_key_values`."
        },
    )
    neftune_noise_alpha: float | None = field(
        default=None,
        metadata={
            "help": "If not None, activates NEFTune noise embeddings. Can drastically improve performance for instruction fine-tuning. Typical range: [5.0, 15.0]."
        },
    )
    torch_empty_cache_steps: int | None = field(
        default=None,
        metadata={
            "help": "Number of steps to wait before calling `torch.<device>.empty_cache()`. Helps avoid CUDA OOM at a cost of ~10%% slower performance. If None, cache will not be emptied."
        },
    )
    auto_find_batch_size: bool = field(
        default=False,
        metadata={
            "help": "Whether to find a batch size that will fit into memory automatically through exponential decay, avoiding CUDA Out-of-Memory errors."
        },
    )

    # --- Logging & Monitoring ---
    logging_strategy: IntervalStrategy | str = field(
        default="steps",
        metadata={"help": "The logging strategy to adopt during training. Options: 'no', 'epoch', 'steps'."},
    )
    logging_steps: float = field(
        default=500,
        metadata={
            "help": (
                "Log every X updates steps. Should be an integer or a float in range `[0,1)`. "
                "If smaller than 1, will be interpreted as ratio of total training steps."
            )
        },
    )
    logging_first_step: bool = field(
        default=False, metadata={"help": "Whether to log the first `global_step` or not."}
    )
    log_on_each_node: bool = field(
        default=True,
        metadata={
            "help": (
                "When doing a multinode distributed training, whether to log once per node or just once on the main"
                " node."
            )
        },
    )
    logging_nan_inf_filter: bool = field(
        default=True,
        metadata={
            "help": "Filter out NaN and Inf losses when logging. Does not affect gradient computation, only logging."
        },
    )
    include_num_input_tokens_seen: str | bool = field(
        default="no",
        metadata={
            "help": (
                "Whether to track the number of input tokens seen. "
                "Must be one of [`all`, `non_padding`, `no`] or a boolean value which map to `all` or `no`"
            )
        },
    )

    # --- Log Levels ---
    log_level: str = field(
        default="passive",
        metadata={
            "help": "Logging level for the main process. Options: 'debug', 'info', 'warning', 'error', 'critical', 'passive'.",
            "choices": trainer_log_levels.keys(),
        },
    )
    log_level_replica: str = field(
        default="warning",
        metadata={
            "help": "Logging level for replica processes in distributed training. Same options as `log_level`.",
            "choices": trainer_log_levels.keys(),
        },
    )
    disable_tqdm: bool | None = field(
        default=None,
        metadata={"help": "Disable tqdm progress bars. Defaults to True if log_level is warning or lower."},
    )

    # --- Experiment Tracking ---
    report_to: None | str | list[str] = field(
        default="none",
        metadata={
            "help": "The list of integrations to report the results and logs to. Use 'all' for all installed integrations, 'none' for no integrations."
        },
    )
    run_name: str | None = field(
        default=None,
        metadata={
            "help": (
                "An optional descriptor for the run. Notably used for trackio, wandb, mlflow comet and swanlab "
                "logging."
            )
        },
    )
    project: str = field(
        default="huggingface",
        metadata={"help": "The name of the project to use for logging. Currently, only used by Trackio."},
    )
    trackio_space_id: str | None = field(
        default=None,
        metadata={
            "help": (
                "Hugging Face Space id for live Gradio-based Trackio logging (read/write Bucket access). Use "
                "'username/reponame', 'orgname/reponame', or 'reponame' (current user's namespace). None: log only "
                "locally, no Space. Prefer trackio_static_space_id for stable post-training dashboard links. Public "
                "unless hub_private_repo=True or org default."
            )
        },
    )
    trackio_bucket_id: str | None = field(
        default=None,
        metadata={"help": "Optional HF Bucket id when using a Trackio Space; if unset, Trackio picks a default."},
    )
    trackio_static_space_id: str | None | Literal[False] = field(
        default=None,
        metadata={
            "help": (
                "Static read-only Trackio Space over the Bucket (stable model-card links). False: no static sync on Hub "
                "push and no freeze after training. None/str: allow static Space; for local-only logging, Hub push runs "
                "sync(static); after training, freeze runs only if trackio_space_id was set (Gradio Space). str sets "
                "explicit static Space id. Public unless hub_private_repo=True or org default."
            )
        },
    )

    # --- Evaluation ---
    eval_strategy: IntervalStrategy | str = field(
        default="no",
        metadata={"help": "When to run evaluation. Options: 'no', 'steps', 'epoch'."},
    )
    eval_steps: float | None = field(
        default=None,
        metadata={
            "help": (
                "Number of update steps between evaluations if `eval_strategy='steps'`. Defaults to `logging_steps` if not set."
                " Should be an integer or a float in range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps."
            )
        },
    )
    eval_delay: float = field(
        default=0,
        metadata={
            "help": (
                "Number of epochs or steps to wait for before the first evaluation can be performed, depending on the"
                " eval_strategy."
            )
        },
    )
    per_device_eval_batch_size: int = field(
        default=8, metadata={"help": "The batch size per device (GPU/TPU core/CPU) for evaluation."}
    )
    prediction_loss_only: bool = field(
        default=False,
        metadata={"help": "When performing evaluation and generating predictions, only returns the loss."},
    )
    eval_on_start: bool = field(
        default=False,
        metadata={
            "help": "Whether to run through the entire `evaluation` step at the very beginning of training as a sanity check."
        },
    )
    eval_do_concat_batches: bool = field(
        default=True,
        metadata={
            "help": "Whether to recursively concat inputs/losses/labels/predictions across batches. If `False`, will instead store them as lists, with each batch kept separate."
        },
    )
    eval_use_gather_object: bool = field(
        default=False,
        metadata={
            "help": "Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices."
        },
    )
    eval_accumulation_steps: int | None = field(
        default=None,
        metadata={
            "help": "Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. If unset, predictions are accumulated on the accelerator before being moved to the CPU."
        },
    )

    # --- Metrics ---
    include_for_metrics: list[str] = field(
        default_factory=list,
        metadata={"help": "Include additional data in the `compute_metrics` function. Options: 'inputs', 'loss'."},
    )
    batch_eval_metrics: bool = field(
        default=False,
        metadata={"help": "Break eval metrics calculation into batches to save memory."},
    )

    # --- Checkpointing & Saving ---
    save_only_model: bool = field(
        default=False,
        metadata={
            "help": "Save only model weights, not optimizer/scheduler/RNG state. Prevents resuming training from checkpoint."
        },
    )
    save_strategy: SaveStrategy | str = field(
        default="steps",
        metadata={
            "help": "The checkpoint save strategy to adopt during training. Options: 'no', 'epoch', 'steps', 'best'."
        },
    )
    save_steps: float = field(
        default=500,
        metadata={
            "help": (
                "Save checkpoint every X updates steps. Should be an integer or a float in range `[0,1)`. "
                "If smaller than 1, will be interpreted as ratio of total training steps."
            )
        },
    )
    save_on_each_node: bool = field(
        default=False,
        metadata={
            "help": (
                "When doing multi-node distributed training, whether to save models and checkpoints on each node, or"
                " only on the main one"
            )
        },
    )
    save_total_limit: int | None = field(
        default=None,
        metadata={
            "help": "Maximum number of checkpoints to keep. Deletes older checkpoints in `output_dir`. The best checkpoint is always retained when `load_best_model_at_end=True`."
        },
    )
    enable_jit_checkpoint: bool = field(
        default=False,
        metadata={
            "help": "Enable JIT checkpointing on SIGTERM signal for graceful termination on preemptible workloads. Configure your orchestrator's graceful shutdown period accordingly."
        },
    )

    # --- Hub Integration ---
    push_to_hub: bool = field(
        default=False, metadata={"help": "Whether or not to push the model to the Hub every time the model is saved."}
    )
    hub_token: str | None = field(
        default=None,
        metadata={
            "help": "The token to use to push the model to the Hub. Defaults to the token from `hf auth login`."
        },
    )
    hub_private_repo: bool | None = field(
        default=None,
        metadata={
            "help": "Whether to make the repo private. If `None` (default), the repo will be public unless the "
            "organization's default is private. This value is ignored if the repo already exists. If reporting to "
            "Trackio Spaces created or synced (including on Hub push when `trackio_space_id` is None) use the same "
            "logic for whether the Space is private."
        },
    )
    hub_model_id: str | None = field(
        default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."}
    )
    hub_strategy: HubStrategy | str = field(
        default="every_save",
        metadata={
            "help": "Defines what and when to push to Hub. Options: 'end', 'every_save', 'checkpoint', 'all_checkpoints'."
        },
    )
    hub_always_push: bool = field(
        default=False,
        metadata={"help": "Unless `True`, the Trainer will skip pushes if the previous one wasn't finished yet."},
    )
    hub_revision: str | None = field(
        default=None,
        metadata={
            "help": "The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash."
        },
    )

    # --- Best Model Tracking ---
    load_best_model_at_end: bool = field(
        default=False,
        metadata={"help": "Load the best checkpoint at the end of training. Requires `eval_strategy` to be set."},
    )
    metric_for_best_model: str | None = field(
        default=None,
        metadata={
            "help": "Metric to use for comparing models when `load_best_model_at_end=True`. Defaults to 'loss'."
        },
    )
    greater_is_better: bool | None = field(
        default=None,
        metadata={"help": "Whether higher metric values are better. Defaults based on `metric_for_best_model`."},
    )

    # --- Resuming Training ---
    ignore_data_skip: bool = field(
        default=False,
        metadata={
            "help": "When resuming training, skip fast-forwarding through the dataset to reach the previous state. If True, training starts from the beginning of the dataset."
        },
    )
    restore_callback_states_from_checkpoint: bool = field(
        default=False,
        metadata={
            "help": "Whether to restore the callback states from the checkpoint. If `True`, will override callbacks passed to the `Trainer` if they exist in the checkpoint."
        },
    )

    # --- Reproducibility ---
    full_determinism: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to call enable_full_determinism instead of set_seed for reproducibility in distributed"
                " training. Important: this will negatively impact the performance, so only use it for debugging."
            )
        },
    )
    seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."})
    data_seed: int | None = field(
        default=None,
        metadata={"help": "Random seed to be used with data samplers. If not set, uses the same seed as `seed`."},
    )

    # --- Hardware ---
    use_cpu: bool = field(
        default=False,
        metadata={
            "help": "Whether or not to use cpu. If set to False, we will use the available torch device/backend."
        },
    )

    # --- Accelerate ---
    accelerator_config: dict | str | None = field(
        default=None,
        metadata={
            "help": "Configuration for the internal Accelerate integration. Can be a path to a JSON config file or a dict."
        },
    )
    parallelism_config: ParallelismConfig | None = field(
        default=None,
        metadata={"help": "Parallelism configuration for the training run. Requires Accelerate `1.10.1`."},
    )

    # --- Dataloader ---
    dataloader_drop_last: bool = field(
        default=False, metadata={"help": "Drop the last incomplete batch if it is not divisible by the batch size."}
    )
    dataloader_num_workers: int = field(
        default=0,
        metadata={
            "help": (
                "Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded"
                " in the main process."
            )
        },
    )
    dataloader_pin_memory: bool = field(
        default=True, metadata={"help": "Whether or not to pin memory for DataLoader."}
    )
    dataloader_persistent_workers: bool = field(
        default=False,
        metadata={
            "help": "If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will increase RAM usage."
        },
    )
    dataloader_prefetch_factor: int | None = field(
        default=None,
        metadata={
            "help": (
                "Number of batches loaded in advance by each worker. "
                "2 means there will be a total of 2 * num_workers batches prefetched across all workers. "
            )
        },
    )
    remove_unused_columns: bool = field(
        default=True,
        metadata={"help": "Whether or not to automatically remove the columns unused by the model forward method."},
    )
    label_names: list[str] | None = field(
        default=None, metadata={"help": "The list of keys in your dictionary of inputs that correspond to the labels."}
    )
    train_sampling_strategy: str = field(
        default="random",
        metadata={
            "help": "Sampler for training: 'random' (default), 'sequential', or 'group_by_length'.",
            "choices": ["random", "sequential", "group_by_length"],
        },
    )
    length_column_name: str = field(
        default="length",
        metadata={
            "help": "Column name for precomputed lengths. Ignored unless `train_sampling_strategy` is 'group_by_length'."
        },
    )

    # --- DDP ---
    ddp_find_unused_parameters: bool | None = field(
        default=None,
        metadata={
            "help": (
                "When using distributed training, the value of the flag `find_unused_parameters` passed to "
                "`DistributedDataParallel`."
            )
        },
    )
    ddp_bucket_cap_mb: int | None = field(
        default=None,
        metadata={
            "help": (
                "When using distributed training, the value of the flag `bucket_cap_mb` passed to "
                "`DistributedDataParallel`."
            )
        },
    )
    ddp_broadcast_buffers: bool | None = field(
        default=None,
        metadata={
            "help": (
                "When using distributed training, the value of the flag `broadcast_buffers` passed to "
                "`DistributedDataParallel`."
            )
        },
    )
    ddp_static_graph: bool | None = field(
        default=None,
        metadata={
            "help": (
                "When using distributed training, the value of the flag `static_graph` passed to "
                "`DistributedDataParallel`."
            )
        },
    )
    ddp_backend: str | None = field(
        default=None,
        metadata={
            "help": "The backend to use for distributed training. Must be one of 'nccl', 'mpi', 'xccl', 'gloo', 'hccl'.",
            "choices": ["nccl", "gloo", "mpi", "xccl", "hccl", "cncl", "mccl"],
        },
    )
    ddp_timeout: int = field(
        default=1800,
        metadata={"help": "The timeout for `torch.distributed.init_process_group` calls (in seconds)."},
    )

    # --- FSDP ---
    fsdp: list[FSDPOption] | str | None = field(
        default=None,
        metadata={
            "help": "Enable PyTorch FSDP for distributed training. Options: 'full_shard', 'shard_grad_op', 'hybrid_shard', 'hybrid_shard_zero2', 'offload', 'auto_wrap'.",
        },
    )
    fsdp_config: dict[str, Any] | str | None = field(
        default=None,
        metadata={
            "help": (
                "Config to be used with FSDP (Pytorch Fully Sharded Data Parallel). The value is either a "
                "fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`."
            )
        },
    )

    # --- DeepSpeed ---
    deepspeed: dict | str | None = field(
        default=None,
        metadata={"help": "Enable DeepSpeed integration. Value is a path to a JSON config file or a dict."},
    )

    # --- Debugging ---
    debug: str | list[DebugOption] = field(
        default="",
        metadata={
            "help": "Enable one or more debug features. Options: 'underflow_overflow' (detect overflow in model I/O), 'tpu_metrics_debug' (print TPU metrics)."
        },
    )
    skip_memory_metrics: bool = field(
        default=True,
        metadata={
            "help": "Whether to skip adding memory profiler reports to metrics. Skipped by default because it slows down training."
        },
    )

    # --- External Script Flags ---
    do_train: bool = field(
        default=False,
        metadata={
            "help": "Whether to run training. Not directly used by Trainer; intended for training/evaluation scripts."
        },
    )
    do_eval: bool = field(
        default=False,
        metadata={
            "help": "Whether to run evaluation. Not directly used by Trainer; intended for training/evaluation scripts."
        },
    )
    do_predict: bool = field(
        default=False,
        metadata={
            "help": "Whether to run predictions on the test set. Not directly used by Trainer; intended for training/evaluation scripts."
        },
    )
    resume_from_checkpoint: str | None = field(
        default=None,
        metadata={
            "help": "Path to a folder with a valid checkpoint for your model. Not directly used by Trainer; intended for training/evaluation scripts."
        },
    )

    # --- Deprecated / Internal ---
    warmup_ratio: float | None = field(
        default=None,
        metadata={
            "help": "This argument is deprecated and will be removed in v5.2. Use `warmup_steps` instead as it also works with float values."
        },
    )
    logging_dir: str | None = field(
        default=None,
        metadata={
            "help": "Deprecated and will be removed in v5.2. Set env var `TENSORBOARD_LOGGING_DIR` instead. TensorBoard log directory."
        },
    )
    local_rank: int = field(
        default=-1,
        metadata={
            "help": "When using torch.distributed.launch (Deprecated), it will pass `local_rank` in the script, so we need this for the parser. To get the local rank, prefer using the property `local_process_index`"
        },
    )

    def __post_init__(self):
        # ── 1. Defaults & Normalization ──
        if self.output_dir is None:
            self.output_dir = "trainer_output"
            logger.info(
                "No output directory specified, defaulting to 'trainer_output'. "
                "To change this behavior, specify --output_dir when creating TrainingArguments."
            )

        # Parse JSON string dict args from CLI (e.g., '{"key": "value"}').
        # Only parses strings starting with '{'; other strings are treated as file paths.
        for valid_field in self._VALID_DICT_FIELDS:
            passed_value = getattr(self, valid_field)
            if isinstance(passed_value, str) and passed_value.startswith("{"):
                loaded_dict = json.loads(passed_value)
                loaded_dict = _convert_str_dict(loaded_dict)
                setattr(self, valid_field, loaded_dict)

        # Expand ~ in paths so os.makedirs works correctly (#10628)
        if self.output_dir is not None:
            self.output_dir = os.path.expanduser(self.output_dir)

        if self.disable_tqdm is None:
            self.disable_tqdm = logger.getEffectiveLevel() > logging.WARN

        if self.warmup_ratio is not None:
            logger.warning("warmup_ratio is deprecated and will be removed in v5.2. Use `warmup_steps` instead.")
            self.warmup_steps = self.warmup_ratio

        if self.logging_dir is not None:
            logger.warning(
                "`logging_dir` is deprecated and will be removed in v5.2. Please set `TENSORBOARD_LOGGING_DIR` instead."
            )

        if isinstance(self.include_num_input_tokens_seen, bool):
            self.include_num_input_tokens_seen = "all" if self.include_num_input_tokens_seen else "no"

        # ── 2. Enum / Type Conversions ──
        self.eval_strategy = IntervalStrategy(self.eval_strategy)
        self.logging_strategy = IntervalStrategy(self.logging_strategy)
        self.save_strategy = SaveStrategy(self.save_strategy)
        self.hub_strategy = HubStrategy(self.hub_strategy)
        self.lr_scheduler_type = SchedulerType(self.lr_scheduler_type)
        self.optim = OptimizerNames(self.optim)

        if isinstance(self.debug, str):
            self.debug = [DebugOption(s) for s in self.debug.split()]
        elif self.debug is None:
            self.debug = []

        # ── 3. Auto-derived Values ──
        if self.do_eval is False and self.eval_strategy != IntervalStrategy.NO:
            self.do_eval = True

        # Fall back to logging_steps if eval_steps is unset
        if self.eval_strategy == IntervalStrategy.STEPS and (self.eval_steps is None or self.eval_steps == 0):
            if self.logging_steps > 0:
                logger.info(f"using `logging_steps` to initialize `eval_steps` to {self.logging_steps}")
                self.eval_steps = self.logging_steps
            else:
                raise ValueError(
                    f"evaluation strategy {self.eval_strategy} requires either non-zero --eval_steps or"
                    " --logging_steps"
                )

        if (
            self.load_best_model_at_end
            or self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU
            or self.lr_scheduler_type == SchedulerType.GREEDY
        ) and self.metric_for_best_model is None:
            self.metric_for_best_model = "loss"
        if self.greater_is_better is None and self.metric_for_best_model is not None:
            self.greater_is_better = not self.metric_for_best_model.endswith("loss")

        if self.report_to == "all" or self.report_to == ["all"]:
            from .integrations import get_available_reporting_integrations

            self.report_to = get_available_reporting_integrations()
        elif self.report_to == "none" or self.report_to == ["none"]:
            self.report_to = []
        elif not isinstance(self.report_to, list):
            self.report_to = [self.report_to]

        # Auto-enable Kubeflow integration when running inside a Kubeflow TrainJob
        from .integrations import is_kubeflow_available

        if is_kubeflow_available() and "kubeflow" not in self.report_to:
            self.report_to = list(self.report_to) + ["kubeflow"]

        # ── 4. Validation ──
        self._validate_args()

        # ── 5. Mixed Precision ──
        # Read from env first; DeepSpeed may override this later
        self.mixed_precision = os.environ.get("ACCELERATE_MIXED_PRECISION", "no")
        if self.fp16:
            self.mixed_precision = "fp16"
        elif self.bf16:
            self.mixed_precision = "bf16"

        # ── 6. Torch Compile ──
        if (self.torch_compile_mode is not None or self.torch_compile_backend is not None) and not self.torch_compile:
            self.torch_compile = True
        if self.torch_compile and self.torch_compile_backend is None:
            if not self.use_cpu and is_torch_hpu_available():
                self.torch_compile_backend = "hpu_backend"
            else:
                self.torch_compile_backend = "inductor"

        if self.torch_compile:
            # TODO: remove env var fallback once minimum accelerate >= 1.2.0
            if not is_accelerate_available("1.2.0"):
                os.environ["ACCELERATE_DYNAMO_BACKEND"] = self.torch_compile_backend
                if self.torch_compile_mode is not None:
                    os.environ["ACCELERATE_DYNAMO_MODE"] = self.torch_compile_mode

        # ── 7. Accelerator Config (must come before self.device) ──
        if is_accelerate_available():
            if not isinstance(self.accelerator_config, AcceleratorConfig):
                if self.accelerator_config is None:
                    self.accelerator_config = AcceleratorConfig()
                elif isinstance(self.accelerator_config, dict):
                    self.accelerator_config = AcceleratorConfig(**self.accelerator_config)
                # Reject uninstantiated class (e.g. AcceleratorConfig instead of AcceleratorConfig())
                elif isinstance(self.accelerator_config, type):
                    raise NotImplementedError(
                        "Tried passing in a callable to `accelerator_config`, but this is not supported. "
                        "Please pass in a fully constructed `AcceleratorConfig` object instead."
                    )
                else:
                    self.accelerator_config = AcceleratorConfig.from_json_file(self.accelerator_config)
            if self.accelerator_config.split_batches:
                logger.info(
                    "Using `split_batches=True` in `accelerator_config` will override the `per_device_train_batch_size` "
                    "Batches will be split across all processes equally when using `split_batches=True`."
                )

        # ── 8. Device Init ──
        if is_torch_available():
            self.device

        # ── 9. TF32 ──
        if is_torch_available() and self.torch_compile:
            if is_torch_tf32_available():
                if self.tf32 is None and not self.fp16 or self.bf16:
                    device_str = "MUSA" if is_torch_musa_available() else "CUDA"
                    logger.info(
                        f"Setting TF32 in {device_str} backends to speedup torch compile, you won't see any improvement"
                        " otherwise."
                    )
                    enable_tf32(True)
            else:
                logger.warning(
                    "The speedups for torchdynamo mostly come with GPU Ampere or higher and which is not detected here."
                )
        if is_torch_available() and self.tf32 is not None:
            if self.tf32:
                if is_torch_tf32_available():
                    enable_tf32(True)
                else:
                    raise ValueError("--tf32 requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7")
            else:
                if is_torch_tf32_available():
                    enable_tf32(False)
                # TF32 not available, nothing to disable

        # ── 10. Hardware Overrides ──
        if self.use_cpu:
            self.dataloader_pin_memory = False

        # ── 11. FSDP ──
        # Store args only (not the plugin itself) to avoid pickle issues
        self.fsdp_plugin_args = self._process_fsdp_args()

        # ── 12. DeepSpeed (must be last) ──
        self.deepspeed_plugin = None
        if self.deepspeed:
            from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig

            # Leave self.deepspeed unmodified; users may rely on the original value
            self.hf_deepspeed_config = HfTrainerDeepSpeedConfig(self.deepspeed)
            self.hf_deepspeed_config.trainer_config_process(self)

            from accelerate.utils import DeepSpeedPlugin

            self.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.hf_deepspeed_config)
        elif strtobool(os.environ.get("ACCELERATE_USE_DEEPSPEED", "false")):
            from accelerate.utils import DeepSpeedPlugin

            self.deepspeed_plugin = DeepSpeedPlugin()
            self.deepspeed_plugin.set_mixed_precision(self.mixed_precision)
            self.deepspeed_plugin.set_deepspeed_weakref()

    def _validate_args(self):
        """Validate argument combinations and value constraints."""
        if self.torch_empty_cache_steps is not None:
            if not (isinstance(self.torch_empty_cache_steps, int) and self.torch_empty_cache_steps > 0):
                raise ValueError(
                    f"`torch_empty_cache_steps` must be an integer bigger than 0, got {self.torch_empty_cache_steps}."
                )

        # logging_steps must be non-zero when logging_strategy="steps"
        if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps == 0:
            raise ValueError(f"logging strategy {self.logging_strategy} requires non-zero --logging_steps")

        if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps > 1:
            if self.logging_steps != int(self.logging_steps):
                raise ValueError(f"--logging_steps must be an integer if bigger than 1: {self.logging_steps}")
            self.logging_steps = int(self.logging_steps)
        if self.eval_strategy == IntervalStrategy.STEPS and self.eval_steps > 1:
            if self.eval_steps != int(self.eval_steps):
                raise ValueError(f"--eval_steps must be an integer if bigger than 1: {self.eval_steps}")
            self.eval_steps = int(self.eval_steps)
        if self.save_strategy == SaveStrategy.STEPS and self.save_steps > 1:
            if self.save_steps != int(self.save_steps):
                raise ValueError(f"--save_steps must be an integer if bigger than 1: {self.save_steps}")
            self.save_steps = int(self.save_steps)

        # load_best_model_at_end requires compatible save and eval strategies
        if self.load_best_model_at_end and self.save_strategy != SaveStrategy.BEST:
            if self.eval_strategy != self.save_strategy:
                raise ValueError(
                    '--load_best_model_at_end requires the save and eval strategy to match, except when --save_strategy="best", but found\n- Evaluation '
                    f"strategy: {self.eval_strategy}\n- Save strategy: {self.save_strategy}"
                )
            if self.eval_strategy == IntervalStrategy.STEPS and self.save_steps % self.eval_steps != 0:
                if self.eval_steps < 1 or self.save_steps < 1:
                    if not (self.eval_steps < 1 and self.save_steps < 1):
                        raise ValueError(
                            "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation "
                            "steps, which cannot get guaranteed when mixing ratio and absolute steps for save_steps "
                            f"{self.save_steps} and eval_steps {self.eval_steps}."
                        )
                    # Use integer arithmetic to avoid floating point precision issues
                    LARGE_MULTIPLIER = 1_000_000
                    if (self.save_steps * LARGE_MULTIPLIER) % (self.eval_steps * LARGE_MULTIPLIER) != 0:
                        raise ValueError(
                            "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation "
                            f"steps, but found {self.save_steps}, which is not a multiple of {self.eval_steps}."
                        )
                else:
                    raise ValueError(
                        "--load_best_model_at_end requires the saving steps to be a round multiple of the evaluation "
                        f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}."
                    )

        if is_torch_available():
            if self.bf16 or self.bf16_full_eval:
                if not self.use_cpu and not is_torch_bf16_gpu_available() and not is_torch_xla_available():
                    error_message = "Your setup doesn't support bf16/gpu. You need to assign use_cpu if you want to train the model on CPU."
                    if is_torch_cuda_available():
                        error_message += " You need Ampere+ GPU with cuda>=11.0."
                    raise ValueError(error_message)

        if self.fp16 and self.bf16:
            raise ValueError("At most one of fp16 and bf16 can be True, but not both")

        if self.fp16_full_eval and self.bf16_full_eval:
            raise ValueError("At most one of fp16 and bf16 can be True for full eval, but not both")

        if self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU:
            if self.eval_strategy == IntervalStrategy.NO:
                raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires an eval strategy")
            if not is_torch_available():
                raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires torch>=0.2.0")

        if self.lr_scheduler_type == SchedulerType.GREEDY:
            if self.eval_strategy == IntervalStrategy.NO:
                raise ValueError("lr_scheduler_type greedy requires an eval strategy")

        if self.warmup_steps < 0:
            raise ValueError("warmup_steps must be an integer or a float")

        if self.dataloader_num_workers == 0 and self.dataloader_prefetch_factor is not None:
            raise ValueError(
                "--dataloader_prefetch_factor can only be set when data is loaded in a different process, i.e."
                " when --dataloader_num_workers > 0."
            )

    def __str__(self):
        self_as_dict = asdict(self)

        self_as_dict = {k: f"<{k.upper()}>" if k.endswith("_token") else v for k, v in self_as_dict.items()}

        attrs_as_str = [f"{k}={v},\n" for k, v in sorted(self_as_dict.items())]
        return f"{self.__class__.__name__}(\n{''.join(attrs_as_str)})"

    __repr__ = __str__

    @property
    def train_batch_size(self) -> int:
        """
        The actual batch size for training.
        """
        train_batch_size = self.per_device_train_batch_size * max(1, self.n_gpu)
        return train_batch_size

    @property
    def eval_batch_size(self) -> int:
        """
        The actual batch size for evaluation.
        """
        eval_batch_size = self.per_device_eval_batch_size * max(1, self.n_gpu)
        return eval_batch_size

    @property
    def ddp_timeout_delta(self) -> timedelta:
        """
        The actual timeout for torch.distributed.init_process_group since it expects a timedelta variable.
        """
        return timedelta(seconds=self.ddp_timeout)

    @cached_property
    def _setup_devices(self) -> "torch.device":
        requires_backends(self, ["torch"])
        logger.info("PyTorch: setting up devices")
        if not is_sagemaker_mp_enabled():
            if not is_accelerate_available():
                raise ImportError(
                    f"Using the `Trainer` with `PyTorch` requires `accelerate>={ACCELERATE_MIN_VERSION}`: "
                    f"Please run `pip install transformers[torch]` or `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`"
                )
        # Build kwargs for PartialState; actual init happens below
        accelerator_state_kwargs: dict[str, Any] = {"enabled": True, "use_configured_state": False}
        if isinstance(self.accelerator_config, AcceleratorConfig):
            accelerator_state_kwargs["use_configured_state"] = self.accelerator_config.pop(
                "use_configured_state", False
            )
        if accelerator_state_kwargs["use_configured_state"]:
            if PartialState._shared_state == {}:
                raise ValueError(
                    "Passing `'use_configured_state':True` to the AcceleratorConfig requires a pre-configured "
                    "`AcceleratorState` or `PartialState` to be defined before calling `TrainingArguments`. "
                )
            self.distributed_state = PartialState(cpu=self.use_cpu)
            if self.deepspeed and self.distributed_state.distributed_type != DistributedType.DEEPSPEED:
                raise RuntimeError(
                    "Tried to use an already configured `Accelerator` or `PartialState` that was not initialized for DeepSpeed, "
                    "but also passed in a `deepspeed` configuration to the `TrainingArguments`. Please set "
                    "`use_configured_state:False` instead or setup your `Accelerator` or `PartialState` properly."
                )
        else:
            AcceleratorState._reset_state(reset_partial_state=True)
            self.distributed_state = None

        self._n_gpu = 1
        if self.use_cpu or strtobool(os.environ.get("ACCELERATE_USE_CPU", "False")):
            accelerator_state_kwargs["cpu"] = True
            accelerator_state_kwargs["backend"] = self.ddp_backend
            self._n_gpu = 0
        elif is_sagemaker_mp_enabled():
            accelerator_state_kwargs["enabled"] = False
            device = torch.device("cuda", smp.local_rank())
            torch.cuda.set_device(device)
        elif is_sagemaker_dp_enabled():
            accelerator_state_kwargs["_use_sagemaker_dp"] = True
        elif self.deepspeed:
            accelerator_state_kwargs["use_deepspeed"] = True
            accelerator_state_kwargs["timeout"] = timedelta(seconds=self.ddp_timeout)
        else:
            accelerator_state_kwargs["backend"] = self.ddp_backend
            accelerator_state_kwargs["timeout"] = timedelta(seconds=self.ddp_timeout)

        # Initialize PartialState with the accumulated kwargs
        if accelerator_state_kwargs.pop("enabled", False) and not accelerator_state_kwargs.pop(
            "use_configured_state", False
        ):
            # Temporarily set env var so Accelerate detects DeepSpeed
            use_deepspeed = accelerator_state_kwargs.pop("use_deepspeed", False)
            if use_deepspeed:
                os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
            self.distributed_state = PartialState(**accelerator_state_kwargs)
            if use_deepspeed:
                del os.environ["ACCELERATE_USE_DEEPSPEED"]
        if not is_sagemaker_mp_enabled():
            device = self.distributed_state.device
        if dist.is_available() and dist.is_initialized() and self.parallel_mode != ParallelMode.DISTRIBUTED:
            logger.warning(
                "torch.distributed process group is initialized, but parallel_mode != ParallelMode.DISTRIBUTED. "
                "In order to use Torch DDP, launch your script with `python -m torch.distributed.launch"
            )

        if is_torch_xla_available():
            device = self.distributed_state.device
            self._n_gpu = 0
        elif is_sagemaker_dp_enabled() or is_sagemaker_mp_enabled():
            pass  # _n_gpu already set above
        elif self.distributed_state.distributed_type == DistributedType.NO:
            if self.use_cpu:
                device = torch.device("cpu")
            elif is_torch_mps_available():
                device = torch.device("mps")
            elif is_torch_xpu_available():
                device = torch.device("xpu:0")
                torch.xpu.set_device(device)
            elif is_torch_mlu_available():
                device = torch.device("mlu:0")
                torch.mlu.set_device(device)
            elif is_torch_musa_available():
                device = torch.device("musa:0")
                torch.musa.set_device(device)
            elif is_torch_npu_available():
                device = torch.device("npu:0")
                torch.npu.set_device(device)
            elif is_torch_hpu_available():
                device = torch.device("hpu:0")
                torch.hpu.set_device(device)
            elif is_torch_neuron_available():
                device = torch.device("neuron:0")
                torch.neuron.set_device(device)
            else:
                # Default to cuda:0 (respects CUDA_VISIBLE_DEVICES); nn.DataParallel handles n_gpu > 1
                device = torch.device(
                    "cuda:0" if torch.cuda.is_available() else os.environ.get("ACCELERATE_TORCH_DEVICE", "cpu")
                )
                # _n_gpu may not have been set yet if _setup_devices is called early
                self._n_gpu = torch.cuda.device_count()
                if device.type == "cuda":
                    torch.cuda.set_device(device)
        return device

    @property
    def device(self) -> "torch.device":
        """
        The device used by this process.
        """
        requires_backends(self, ["torch"])
        return self._setup_devices

    @property
    def n_gpu(self):
        """
        The number of GPUs used by this process.

        Note:
            This will only be greater than one when you have multiple GPUs available but are not using distributed
            training. For distributed training, it will always be 1.
        """
        requires_backends(self, ["torch"])
        # Ensure _setup_devices has been called
        if not hasattr(self, "_n_gpu"):
            _ = self._setup_devices
        return self._n_gpu

    @property
    def parallel_mode(self):
        """
        The current mode used for parallelism if multiple GPUs/TPU cores are available. One of:

        - `ParallelMode.NOT_PARALLEL`: no parallelism (CPU or one GPU).
        - `ParallelMode.NOT_DISTRIBUTED`: several GPUs in one single process (uses `torch.nn.DataParallel`).
        - `ParallelMode.DISTRIBUTED`: several GPUs, each having its own process (uses
          `torch.nn.DistributedDataParallel`).
        - `ParallelMode.TPU`: several TPU cores.
        """
        requires_backends(self, ["torch"])
        if is_torch_xla_available():
            return ParallelMode.TPU
        elif is_sagemaker_mp_enabled():
            return ParallelMode.SAGEMAKER_MODEL_PARALLEL
        elif is_sagemaker_dp_enabled():
            return ParallelMode.SAGEMAKER_DATA_PARALLEL
        elif self.distributed_state is not None and self.distributed_state.distributed_type != DistributedType.NO:
            return ParallelMode.DISTRIBUTED
        elif self.n_gpu > 1:
            return ParallelMode.NOT_DISTRIBUTED
        else:
            return ParallelMode.NOT_PARALLEL

    @property
    def world_size(self):
        """
        The number of processes used in parallel.
        """
        requires_backends(self, ["torch"])
        if self.distributed_state is not None:
            return self.distributed_state.num_processes
        elif is_sagemaker_mp_enabled():
            return smp.dp_size() if not smp.state.cfg.prescaled_batch else smp.rdp_size()
        return 1

    @property
    def process_index(self):
        """
        The index of the current process used.
        """
        requires_backends(self, ["torch"])
        if self.distributed_state is not None:
            return self.distributed_state.process_index
        elif is_sagemaker_mp_enabled():
            return smp.dp_rank() if not smp.state.cfg.prescaled_batch else smp.rdp_rank()
        return 0

    @property
    def local_process_index(self):
        """
        The index of the local process used.
        """
        requires_backends(self, ["torch"])

        if self.distributed_state is not None:
            return self.distributed_state.local_process_index
        elif is_sagemaker_mp_enabled():
            return smp.local_rank()
        return 0

    @property
    def should_log(self):
        """
        Whether or not the current process should produce log.
        """
        if self.log_on_each_node:
            return self.local_process_index == 0
        else:
            if is_sagemaker_mp_enabled():
                return smp.rank() == 0
            else:
                return self.process_index == 0

    @property
    def should_save(self):
        """
        Whether or not the current process should write to disk, e.g., to save models and checkpoints.
        """
        if self.save_on_each_node:
            return self.local_process_index == 0
        else:
            if is_sagemaker_mp_enabled():
                return smp.rank() == 0
            else:
                return self.process_index == 0

    def get_process_log_level(self):
        """
        Returns the log level to be used depending on whether this process is the main process of node 0, main process
        of node non-0, or a non-main process.

        For the main process the log level defaults to the logging level set (`logging.WARNING` if you didn't do
        anything) unless overridden by `log_level` argument.

        For the replica processes the log level defaults to `logging.WARNING` unless overridden by `log_level_replica`
        argument.

        The choice between the main and replica process settings is made according to the return value of `should_log`.
        """

        # convert to int
        log_level = trainer_log_levels[self.log_level]
        log_level_replica = trainer_log_levels[self.log_level_replica]

        log_level_main_node = logging.get_verbosity() if log_level == -1 else log_level
        log_level_replica_node = logging.get_verbosity() if log_level_replica == -1 else log_level_replica
        return log_level_main_node if self.should_log else log_level_replica_node

    @property
    def place_model_on_device(self) -> bool | None:
        """
        Can be subclassed and overridden for some specific integrations.
        """
        return None

    @property
    def _no_sync_in_gradient_accumulation(self):
        """
        Whether or not to use no_sync for the gradients when doing gradient accumulation.
        """
        return not (
            self.deepspeed or is_sagemaker_dp_enabled() or is_sagemaker_mp_enabled() or is_torch_neuroncore_available()
        )

    @contextlib.contextmanager
    def main_process_first(self, local=True, desc="work"):
        """
        A context manager for torch distributed environment where on needs to do something on the main process, while
        blocking replicas, and when it's finished releasing the replicas.

        One such use is for `datasets`'s `map` feature which to be efficient should be run once on the main process,
        which upon completion saves a cached version of results and which then automatically gets loaded by the
        replicas.

        Args:
            local (`bool`, *optional*, defaults to `True`):
                if `True` first means process of rank 0 of each node if `False` first means process of rank 0 of node
                rank 0 In multi-node environment with a shared filesystem you most likely will want to use
                `local=False` so that only the main process of the first node will do the processing. If however, the
                filesystem is not shared, then the main process of each node will need to do the processing, which is
                the default behavior.
            desc (`str`, *optional*, defaults to `"work"`):
                a work description to be used in debug logs

        """
        if is_torch_available() and self.world_size > 1:
            main_process_desc = "main local process" if local else "main process"
            if self.distributed_state is not None:
                is_main_process = (
                    self.distributed_state.is_local_main_process if local else self.distributed_state.is_main_process
                )
            elif is_sagemaker_mp_enabled():
                is_main_process = smp.rank() == 0

            try:
                if not is_main_process:
                    # tell all replicas to wait
                    logger.debug(f"{self.process_index}: waiting for the {main_process_desc} to perform {desc}")

                    if is_torch_xla_available():
                        xm.rendezvous(desc)
                    else:
                        dist.barrier()
                yield
            finally:
                if is_main_process:
                    # the wait is over
                    logger.debug(f"{self.process_index}: {main_process_desc} completed {desc}, releasing all replicas")
                    if is_torch_xla_available():
                        xm.rendezvous(desc)
                    else:
                        dist.barrier()
        else:
            yield

    def get_warmup_steps(self, num_training_steps: int):
        """
        Get number of steps used for a linear warmup.
        """
        warmup_steps = (
            int(self.warmup_steps) if self.warmup_steps >= 1 else math.ceil(num_training_steps * self.warmup_steps)
        )
        return warmup_steps

    def _dict_dtype_to_str(self, d: dict[str, Any]) -> None:
        """
        Checks whether the passed dictionary and its nested dicts have a *dtype* key and if it's not None,
        converts torch.dtype to a string of just the type. For example, `torch.float32` get converted into *"float32"*
        string, which can then be stored in the json format.
        """
        if d.get("dtype") is not None and not isinstance(d["dtype"], str):
            d["dtype"] = str(d["dtype"]).split(".")[1]
        for value in d.values():
            if isinstance(value, dict):
                self._dict_dtype_to_str(value)

    def to_dict(self):
        """
        Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates
        the token values by removing their value.
        """
        # Exclude non-init fields (they aren't user-facing config)
        d = {field.name: getattr(self, field.name) for field in fields(self) if field.init}

        for k, v in d.items():
            if isinstance(v, Enum):
                d[k] = v.value
            if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum):
                d[k] = [x.value for x in v]
            if k.endswith("_token"):
                d[k] = f"<{k.upper()}>"
            # Serialize AcceleratorConfig to dict
            if is_accelerate_available() and isinstance(v, AcceleratorConfig):
                d[k] = v.to_dict()
            # Serialize quantization_config if nested inside model_init_kwargs
            if k == "model_init_kwargs" and isinstance(v, dict) and "quantization_config" in v:
                quantization_config = v.get("quantization_config")
                if quantization_config and not isinstance(quantization_config, dict):
                    d[k]["quantization_config"] = quantization_config.to_dict()
            if k == "parallelism_config" and v is not None:
                d[k] = v.to_json()

        self._dict_dtype_to_str(d)

        return d

    def to_json_string(self):
        """
        Serializes this instance to a JSON string.
        """
        return json.dumps(self.to_dict(), indent=2)

    def to_sanitized_dict(self) -> dict[str, Any]:
        """
        Sanitized serialization to use with TensorBoard's hparams
        """
        d = self.to_dict()
        d = {**d, "train_batch_size": self.train_batch_size, "eval_batch_size": self.eval_batch_size}

        valid_types = [bool, int, float, str]
        if is_torch_available():
            valid_types.append(torch.Tensor)

        return {k: v if type(v) in valid_types else str(v) for k, v in d.items()}

    # Convenience setters for grouped configuration
    def set_training(
        self,
        learning_rate: float = 5e-5,
        batch_size: int = 8,
        weight_decay: float = 0,
        num_epochs: float = 3,
        max_steps: int = -1,
        gradient_accumulation_steps: int = 1,
        seed: int = 42,
        gradient_checkpointing: bool = False,
    ):
        """
        A method that regroups all basic arguments linked to the training.

        <Tip>

        Calling this method will automatically set `self.do_train` to `True`.

        </Tip>

        Args:
            learning_rate (`float`, *optional*, defaults to 5e-5):
                The initial learning rate for the optimizer.
            batch_size (`int` *optional*, defaults to 8):
                The batch size per device (GPU/TPU core/CPU...) used for training.
            weight_decay (`float`, *optional*, defaults to 0):
                The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights in the
                optimizer.
            num_train_epochs(`float`, *optional*, defaults to 3.0):
                Total number of training epochs to perform (if not an integer, will perform the decimal part percents
                of the last epoch before stopping training).
            max_steps (`int`, *optional*, defaults to -1):
                If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`.
                For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until
                `max_steps` is reached.
            gradient_accumulation_steps (`int`, *optional*, defaults to 1):
                Number of updates steps to accumulate the gradients for, before performing a backward/update pass.

                <Tip warning={true}>

                When using gradient accumulation, one step is counted as one step with backward pass. Therefore,
                logging, evaluation, save will be conducted every `gradient_accumulation_steps * xxx_step` training
                examples.

                </Tip>

            seed (`int`, *optional*, defaults to 42):
                Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use
                the [`~Trainer.model_init`] function to instantiate the model if it has some randomly initialized
                parameters.
            gradient_checkpointing (`bool`, *optional*, defaults to `False`):
                If True, use gradient checkpointing to save memory at the expense of slower backward pass.

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_training(learning_rate=1e-4, batch_size=32)
        >>> args.learning_rate
        1e-4
        ```
        """
        self.do_train = True
        self.learning_rate = learning_rate
        self.per_device_train_batch_size = batch_size
        self.weight_decay = weight_decay
        self.num_train_epochs = num_epochs
        self.max_steps = max_steps
        self.gradient_accumulation_steps = gradient_accumulation_steps
        self.seed = seed
        self.gradient_checkpointing = gradient_checkpointing
        return self

    def set_evaluate(
        self,
        strategy: str | IntervalStrategy = "no",
        steps: int = 500,
        batch_size: int = 8,
        accumulation_steps: int | None = None,
        delay: float | None = None,
        loss_only: bool = False,
    ):
        """
        A method that regroups all arguments linked to evaluation.

        Args:
            strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"no"`):
                The evaluation strategy to adopt during training. Possible values are:

                    - `"no"`: No evaluation is done during training.
                    - `"steps"`: Evaluation is done (and logged) every `steps`.
                    - `"epoch"`: Evaluation is done at the end of each epoch.

                Setting a `strategy` different from `"no"` will set `self.do_eval` to `True`.
            steps (`int`, *optional*, defaults to 500):
                Number of update steps between two evaluations if `strategy="steps"`.
            batch_size (`int` *optional*, defaults to 8):
                The batch size per device (GPU/TPU core/CPU...) used for evaluation.
            accumulation_steps (`int`, *optional*):
                Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU.
                If left unset, the whole predictions are accumulated on GPU/TPU before being moved to the CPU (faster
                but requires more memory).
            delay (`float`, *optional*):
                Number of epochs or steps to wait for before the first evaluation can be performed, depending on the
                eval_strategy.
            loss_only (`bool`, *optional*, defaults to `False`):
                Ignores all outputs except the loss.

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_evaluate(strategy="steps", steps=100)
        >>> args.eval_steps
        100
        ```
        """
        self.eval_strategy = IntervalStrategy(strategy)
        if self.eval_strategy == IntervalStrategy.STEPS and steps == 0:
            raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.")
        self.do_eval = self.eval_strategy != IntervalStrategy.NO
        self.eval_steps = steps
        self.per_device_eval_batch_size = batch_size
        self.eval_accumulation_steps = accumulation_steps
        self.eval_delay = delay
        self.prediction_loss_only = loss_only
        return self

    def set_testing(
        self,
        batch_size: int = 8,
        loss_only: bool = False,
    ):
        """
        A method that regroups all basic arguments linked to testing on a held-out dataset.

        <Tip>

        Calling this method will automatically set `self.do_predict` to `True`.

        </Tip>

        Args:
            batch_size (`int` *optional*, defaults to 8):
                The batch size per device (GPU/TPU core/CPU...) used for testing.
            loss_only (`bool`, *optional*, defaults to `False`):
                Ignores all outputs except the loss.

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_testing(batch_size=32)
        >>> args.per_device_eval_batch_size
        32
        ```
        """
        self.do_predict = True
        self.per_device_eval_batch_size = batch_size
        self.prediction_loss_only = loss_only
        return self

    def set_save(
        self,
        strategy: str | IntervalStrategy = "steps",
        steps: int = 500,
        total_limit: int | None = None,
        on_each_node: bool = False,
    ):
        """
        A method that regroups all arguments linked to checkpoint saving.

        Args:
            strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`):
                The checkpoint save strategy to adopt during training. Possible values are:

                    - `"no"`: No save is done during training.
                    - `"epoch"`: Save is done at the end of each epoch.
                    - `"steps"`: Save is done every `save_steps`.

            steps (`int`, *optional*, defaults to 500):
                Number of updates steps before two checkpoint saves if `strategy="steps"`.
            total_limit (`int`, *optional*):
                If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in
                `output_dir`.
            on_each_node (`bool`, *optional*, defaults to `False`):
                When doing multi-node distributed training, whether to save models and checkpoints on each node, or
                only on the main one.

                This should not be activated when the different nodes use the same storage as the files will be saved
                with the same names for each node.

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_save(strategy="steps", steps=100)
        >>> args.save_steps
        100
        ```
        """
        self.save_strategy = SaveStrategy(strategy)
        if self.save_strategy == SaveStrategy.STEPS and steps == 0:
            raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.")
        self.save_steps = steps
        self.save_total_limit = total_limit
        self.save_on_each_node = on_each_node
        return self

    def set_logging(
        self,
        strategy: str | IntervalStrategy = "steps",
        steps: int = 500,
        report_to: str | list[str] = "none",
        level: str = "passive",
        first_step: bool = False,
        nan_inf_filter: bool = False,
        on_each_node: bool = False,
        replica_level: str = "passive",
    ):
        """
        A method that regroups all arguments linked to logging.

        Args:
            strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`):
                The logging strategy to adopt during training. Possible values are:

                    - `"no"`: No logging is done during training.
                    - `"epoch"`: Logging is done at the end of each epoch.
                    - `"steps"`: Logging is done every `logging_steps`.

            steps (`int`, *optional*, defaults to 500):
                Number of update steps between two logs if `strategy="steps"`.
            level (`str`, *optional*, defaults to `"passive"`):
                Logger log level to use on the main process. Possible choices are the log levels as strings: `"debug"`,
                `"info"`, `"warning"`, `"error"` and `"critical"`, plus a `"passive"` level which doesn't set anything
                and lets the application set the level.
            report_to (`str` or `list[str]`, *optional*, defaults to `"none"`):
                The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`,
                `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`,
                `"swanlab"`, `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations
                installed, `"none"` for no integrations.
            first_step (`bool`, *optional*, defaults to `False`):
                Whether to log and evaluate the first `global_step` or not.
            nan_inf_filter (`bool`, *optional*, defaults to `True`):
                Whether to filter `nan` and `inf` losses for logging. If set to `True` the loss of every step that is
                `nan` or `inf` is filtered and the average loss of the current logging window is taken instead.

                <Tip>

                `nan_inf_filter` only influences the logging of loss values, it does not change the behavior the
                gradient is computed or applied to the model.

                </Tip>

            on_each_node (`bool`, *optional*, defaults to `True`):
                In multinode distributed training, whether to log using `log_level` once per node, or only on the main
                node.
            replica_level (`str`, *optional*, defaults to `"passive"`):
                Logger log level to use on replicas. Same choices as `log_level`

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_logging(strategy="steps", steps=100)
        >>> args.logging_steps
        100
        ```
        """
        self.logging_strategy = IntervalStrategy(strategy)
        if self.logging_strategy == IntervalStrategy.STEPS and steps == 0:
            raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.")
        self.logging_steps = steps
        self.report_to = report_to
        self.log_level = level
        self.logging_first_step = first_step
        self.logging_nan_inf_filter = nan_inf_filter
        self.log_on_each_node = on_each_node
        self.log_level_replica = replica_level
        return self

    def set_push_to_hub(
        self,
        model_id: str,
        strategy: str | HubStrategy = "every_save",
        token: str | None = None,
        private_repo: bool | None = None,
        always_push: bool = False,
        revision: str | None = None,
    ):
        """
        A method that regroups all arguments linked to synchronizing checkpoints with the Hub.

        <Tip>

        Calling this method will set `self.push_to_hub` to `True`, which means the `output_dir` will begin a git
        directory synced with the repo (determined by `model_id`) and the content will be pushed each time a save is
        triggered (depending on your `self.save_strategy`). Calling [`~Trainer.save_model`] will also trigger a push.

        </Tip>

        Args:
            model_id (`str`):
                The name of the repository to keep in sync with the local *output_dir*. It can be a simple model ID in
                which case the model will be pushed in your namespace. Otherwise it should be the whole repository
                name, for instance `"user_name/model"`, which allows you to push to an organization you are a member of
                with `"organization_name/model"`.
            strategy (`str` or [`~trainer_utils.HubStrategy`], *optional*, defaults to `"every_save"`):
                Defines the scope of what is pushed to the Hub and when. Possible values are:

                - `"end"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`]) and a
                draft of a model card when the [`~Trainer.save_model`] method is called.
                - `"every_save"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`])
                  and
                a draft of a model card each time there is a model save. The pushes are asynchronous to not block
                training, and in case the save are very frequent, a new push is only attempted if the previous one is
                finished. A last push is made with the final model at the end of training.
                - `"checkpoint"`: like `"every_save"` but the latest checkpoint is also pushed in a subfolder named
                last-checkpoint, allowing you to resume training easily with
                `trainer.train(resume_from_checkpoint="last-checkpoint")`.
                - `"all_checkpoints"`: like `"checkpoint"` but all checkpoints are pushed like they appear in the
                  output
                folder (so you will get one checkpoint folder per folder in your final repository)

            token (`str`, *optional*):
                The token to use to push the model to the Hub. Will default to the token in the cache folder obtained
                with `hf auth login`.
            private_repo (`bool`, *optional*, defaults to `False`):
                Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.
            always_push (`bool`, *optional*, defaults to `False`):
                Unless this is `True`, the `Trainer` will skip pushing a checkpoint when the previous push is not
                finished.
            revision (`str`, *optional*):
                The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash.

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_push_to_hub("me/awesome-model")
        >>> args.hub_model_id
        'me/awesome-model'
        ```
        """
        self.push_to_hub = True
        self.hub_model_id = model_id
        self.hub_strategy = HubStrategy(strategy)
        self.hub_token = token
        self.hub_private_repo = private_repo
        self.hub_always_push = always_push
        self.hub_revision = revision
        return self

    def set_optimizer(
        self,
        name: str | OptimizerNames = "adamw_torch",
        learning_rate: float = 5e-5,
        weight_decay: float = 0,
        beta1: float = 0.9,
        beta2: float = 0.999,
        epsilon: float = 1e-8,
        args: str | None = None,
    ):
        """
        A method that regroups all arguments linked to the optimizer and its hyperparameters.

        Args:
            name (`str` or [`training_args.OptimizerNames`], *optional*, defaults to `"adamw_torch"`):
                The optimizer to use: `"adamw_torch"`, `"adamw_torch_fused"`, `"adamw_apex_fused"`,
                `"adamw_anyprecision"` or `"adafactor"`.
            learning_rate (`float`, *optional*, defaults to 5e-5):
                The initial learning rate.
            weight_decay (`float`, *optional*, defaults to 0):
                The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights.
            beta1 (`float`, *optional*, defaults to 0.9):
                The beta1 hyperparameter for the adam optimizer or its variants.
            beta2 (`float`, *optional*, defaults to 0.999):
                The beta2 hyperparameter for the adam optimizer or its variants.
            epsilon (`float`, *optional*, defaults to 1e-8):
                The epsilon hyperparameter for the adam optimizer or its variants.
            args (`str`, *optional*):
                Optional arguments that are supplied to AnyPrecisionAdamW (only useful when
                `optim="adamw_anyprecision"`).

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_optimizer(name="adamw_torch", beta1=0.8)
        >>> args.optim
        'adamw_torch'
        ```
        """
        self.optim = OptimizerNames(name)
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.adam_beta1 = beta1
        self.adam_beta2 = beta2
        self.adam_epsilon = epsilon
        self.optim_args = args
        return self

    def set_lr_scheduler(
        self,
        name: str | SchedulerType = "linear",
        num_epochs: float = 3.0,
        max_steps: int = -1,
        warmup_steps: float = 0,
        warmup_ratio: float | None = None,
    ):
        """
        A method that regroups all arguments linked to the learning rate scheduler and its hyperparameters.

        Args:
            name (`str` or [`SchedulerType`], *optional*, defaults to `"linear"`):
                The scheduler type to use. See the documentation of [`SchedulerType`] for all possible values.
            num_epochs(`float`, *optional*, defaults to 3.0):
                Total number of training epochs to perform (if not an integer, will perform the decimal part percents
                of the last epoch before stopping training).
            max_steps (`int`, *optional*, defaults to -1):
                If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`.
                For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until
                `max_steps` is reached.
            warmup_steps (`float`, *optional*, defaults to 0):
                Number of steps used for a linear warmup from 0 to `learning_rate`.  Should be an integer or a float in range `[0,1)`.
                If smaller than 1, will be interpreted as ratio of steps used for a linear warmup from 0 to `learning_rate`.

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_lr_scheduler(name="cosine", warmup_steps=0.05)
        >>> args.warmup_steps
        0.05
        ```
        """
        if warmup_ratio is not None:
            logger.warning("warmup_ratio is deprecated and will be removed in v5.2 . Use `warmup_steps` instead.")
            warmup_steps = warmup_ratio

        self.lr_scheduler_type = SchedulerType(name)
        self.num_train_epochs = num_epochs
        self.max_steps = max_steps
        self.warmup_steps = warmup_steps
        return self

    def set_dataloader(
        self,
        train_batch_size: int = 8,
        eval_batch_size: int = 8,
        drop_last: bool = False,
        num_workers: int = 0,
        pin_memory: bool = True,
        persistent_workers: bool = False,
        prefetch_factor: int | None = None,
        auto_find_batch_size: bool = False,
        ignore_data_skip: bool = False,
        sampler_seed: int | None = None,
    ):
        """
        A method that regroups all arguments linked to the dataloaders creation.

        Args:
            drop_last (`bool`, *optional*, defaults to `False`):
                Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch
                size) or not.
            num_workers (`int`, *optional*, defaults to 0):
                Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in
                the main process.
            pin_memory (`bool`, *optional*, defaults to `True`):
                Whether you want to pin memory in data loaders or not. Will default to `True`.
            persistent_workers (`bool`, *optional*, defaults to `False`):
                If True, the data loader will not shut down the worker processes after a dataset has been consumed
                once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training,
                but will increase RAM usage. Will default to `False`.
            prefetch_factor (`int`, *optional*):
                Number of batches loaded in advance by each worker.
                2 means there will be a total of 2 * num_workers batches prefetched across all workers.
            auto_find_batch_size (`bool`, *optional*, defaults to `False`)
                Whether to find a batch size that will fit into memory automatically through exponential decay,
                avoiding CUDA Out-of-Memory errors. Requires accelerate to be installed (`pip install accelerate`)
            ignore_data_skip (`bool`, *optional*, defaults to `False`):
                When resuming training, whether or not to skip the epochs and batches to get the data loading at the
                same stage as in the previous training. If set to `True`, the training will begin faster (as that
                skipping step can take a long time) but will not yield the same results as the interrupted training
                would have.
            sampler_seed (`int`, *optional*):
                Random seed to be used with data samplers. If not set, random generators for data sampling will use the
                same seed as `self.seed`. This can be used to ensure reproducibility of data sampling, independent of
                the model seed.

        Example:

        ```py
        >>> from transformers import TrainingArguments

        >>> args = TrainingArguments("working_dir")
        >>> args = args.set_dataloader(train_batch_size=16, eval_batch_size=64)
        >>> args.per_device_train_batch_size
        16
        ```
        """
        self.per_device_train_batch_size = train_batch_size
        self.per_device_eval_batch_size = eval_batch_size
        self.dataloader_drop_last = drop_last
        self.dataloader_num_workers = num_workers
        self.dataloader_pin_memory = pin_memory
        self.dataloader_persistent_workers = persistent_workers
        self.dataloader_prefetch_factor = prefetch_factor
        self.auto_find_batch_size = auto_find_batch_size
        self.ignore_data_skip = ignore_data_skip
        self.data_seed = sampler_seed
        return self

    def _process_fsdp_args(self):
        if not self.fsdp:
            self.fsdp = []
        elif self.fsdp is True:
            self.fsdp = [FSDPOption.FULL_SHARD]
        elif isinstance(self.fsdp, str):
            self.fsdp = [FSDPOption(s) for s in self.fsdp.split()]

        if self.fsdp == [FSDPOption.OFFLOAD]:
            raise ValueError(
                "`--fsdp offload` can't work on its own. It needs to be added to `--fsdp full_shard` or "
                '`--fsdp shard_grad_op`. For example, `--fsdp "full_shard offload"`.'
            )
        elif FSDPOption.FULL_SHARD in self.fsdp and FSDPOption.SHARD_GRAD_OP in self.fsdp:
            raise ValueError("`--fsdp full_shard` is not compatible with `--fsdp shard_grad_op`.")

        if self.gradient_checkpointing and (
            FSDPOption.FULL_SHARD in self.fsdp or FSDPOption.HYBRID_SHARD in self.fsdp
        ):
            logger.warning(
                "When using FSDP full shard, instead of using `gradient_checkpointing` in TrainingArguments, please"
                " use `activation_checkpointing` in `fsdp_config`. The former introduces a redundant AllGather"
                " operation in backward pass. Reference: https://github.com/huggingface/transformers/issues/30404"
            )

        if self.fsdp_config is None:
            self.fsdp_config = {}

        if isinstance(self.fsdp_config, str):
            if len(self.fsdp) == 0:
                warnings.warn("`--fsdp_config` is useful only when `--fsdp` is specified.")
            with open(self.fsdp_config, encoding="utf-8") as f:
                self.fsdp_config = json.load(f)

        if self.fsdp_config is not None and isinstance(self.fsdp_config, dict):
            for k in list(self.fsdp_config.keys()):
                if k.startswith("fsdp_"):
                    v = self.fsdp_config.pop(k)
                    self.fsdp_config[k[5:]] = v

        self.fsdp_config["min_num_params"] = self.fsdp_config.get("min_num_params", 0)

        # Normalize transformer_layer_cls_to_wrap from string to list
        if isinstance(self.fsdp_config.get("transformer_layer_cls_to_wrap", None), str):
            self.fsdp_config["transformer_layer_cls_to_wrap"] = [self.fsdp_config["transformer_layer_cls_to_wrap"]]

        if len(self.fsdp) == 0 and self.fsdp_config["min_num_params"] > 0:
            warnings.warn("`min_num_params` is useful only when `--fsdp` is specified.")

        if len(self.fsdp) == 0 and self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None:
            warnings.warn("`transformer_layer_cls_to_wrap` is useful only when `--fsdp` is specified.")

        if (
            len(self.fsdp) > 0
            and self.fsdp_config["min_num_params"] > 0
            and self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None
        ):
            raise ValueError("`min_num_params` and `transformer_layer_cls_to_wrap` are mutually exclusive.")
        self.fsdp_config["xla"] = self.fsdp_config.get("xla", False)
        self.fsdp_config["xla_fsdp_v2"] = self.fsdp_config.get("xla_fsdp_v2", False)
        self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False)
        if self.fsdp_config["xla"]:
            if len(self.fsdp) > 0:
                # Copy to avoid mutating the original (needed for JSON serialization)
                self.xla_fsdp_config = self.fsdp_config.get("xla_fsdp_settings", {}).copy()
                # Convert string dtype names to torch.dtype
                if "compute_dtype" in self.xla_fsdp_config:
                    self.xla_fsdp_config["compute_dtype"] = getattr(torch, self.xla_fsdp_config["compute_dtype"])
                if "buffer_dtype" in self.xla_fsdp_config:
                    self.xla_fsdp_config["buffer_dtype"] = getattr(torch, self.xla_fsdp_config["buffer_dtype"])
            else:
                warnings.warn("XLA FSDP can be used only when `--fsdp` is specified.")
        else:
            if self.fsdp_config["xla_fsdp_grad_ckpt"]:
                warnings.warn("`--xla_fsdp_grad_ckpt` is useful only when `--xla` is set to true.")

        # Build kwargs for Accelerate's FSDPPlugin
        fsdp_plugin_args = None
        if len(self.fsdp) > 0 and not self.fsdp_config["xla"]:
            from accelerate.utils.constants import (
                FSDP_AUTO_WRAP_POLICY,
                FSDP_SHARDING_STRATEGY,
            )

            fsdp_plugin_args = {}
            fsdp_sharding = None
            for fsdp_option in self.fsdp:
                if fsdp_option.upper() in FSDP_SHARDING_STRATEGY:
                    fsdp_sharding = fsdp_option
                elif fsdp_option == FSDPOption.OFFLOAD:
                    fsdp_plugin_args["cpu_offload"] = True
                elif fsdp_option == FSDPOption.AUTO_WRAP:
                    fsdp_plugin_args["auto_wrap_policy"] = FSDP_AUTO_WRAP_POLICY[0]
                    if self.fsdp_config["min_num_params"] > 0:
                        fsdp_plugin_args["min_num_params"] = self.fsdp_config["min_num_params"]
                        fsdp_plugin_args["auto_wrap_policy"] = FSDP_AUTO_WRAP_POLICY[1]
                    elif self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None:
                        fsdp_plugin_args["transformer_cls_names_to_wrap"] = ",".join(
                            self.fsdp_config["transformer_layer_cls_to_wrap"]
                        )
            fsdp_version = int(self.fsdp_config.get("version", 1))
            fsdp_plugin_args["fsdp_version"] = fsdp_version
            prefetch_policy = self.fsdp_config.get("backward_prefetch", "NO_PREFETCH")
            if fsdp_version == 2:
                # full_shard → True (reshard after forward), shard_grad_op → False
                default_reshard = fsdp_sharding != "shard_grad_op" if fsdp_sharding else True
                fsdp_plugin_args["reshard_after_forward"] = str_to_bool(
                    str(self.fsdp_config.get("reshard_after_forward", default_reshard)).lower()
                )
            else:
                fsdp_plugin_args["forward_prefetch"] = str_to_bool(
                    str(self.fsdp_config.get("forward_prefetch", "false")).lower()
                )
                fsdp_plugin_args["backward_prefetch"] = prefetch_policy.upper()
                # Pass sharding strategy as reshard_after_forward (accelerate converts it to ShardingStrategy)
                default_reshard = fsdp_sharding.upper() if fsdp_sharding else "FULL_SHARD"
                fsdp_plugin_args["reshard_after_forward"] = str(
                    self.fsdp_config.get("reshard_after_forward", default_reshard)
                ).lower()
                fsdp_plugin_args["use_orig_params"] = str_to_bool(
                    str(self.fsdp_config.get("use_orig_params", "true")).lower()
                )

            sync_module_states = str(self.fsdp_config.get("sync_module_states", "true")).lower()
            cpu_ram_efficient_loading = str(self.fsdp_config.get("cpu_ram_efficient_loading", "false")).lower()
            if sync_module_states == "false" and cpu_ram_efficient_loading == "true":
                # Without sync, non-main processes would have random weights
                raise ValueError('`sync_module_states` must be `"True"` if `cpu_ram_efficient_loading` is `"True"`')

            # Set env var to suppress Accelerate warning and for transformers to read
            fsdp_plugin_args["cpu_ram_efficient_loading"] = str_to_bool(cpu_ram_efficient_loading)
            os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = cpu_ram_efficient_loading

            fsdp_plugin_args["sync_module_states"] = str_to_bool(sync_module_states)

        return fsdp_plugin_args

Configuration class for controlling all aspects of model training with the Trainer. TrainingArguments centralizes all hyperparameters, optimization settings, logging preferences, and infrastructure choices needed for training.

[HfArgumentParser] can turn this class into argparse arguments that can be specified on the command line.

Parameters -----= output_dir (str, optional, defaults to "trainer_output"): The output directory where the model predictions and checkpoints will be written.

Training Duration and Batch Size

per_device_train_batch_size (int, optional, defaults to 8): The batch size per device. The global batch size is computed as: per_device_train_batch_size * number_of_devices in multi-GPU or distributed setups. num_train_epochs(float, optional, defaults to 3.0): Total number of training epochs to perform (if not an integer, will perform the decimal part percents of the last epoch before stopping training). max_steps (int, optional, defaults to -1): Overrides num_train_epochs. If set to a positive number, the total number of training steps to perform. For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until max_steps is reached.

Learning Rate & Scheduler

learning_rate (float, optional, defaults to 5e-5): The initial learning rate for the optimizer. This is typically the peak learning rate when using a scheduler with warmup. lr_scheduler_type (str or [SchedulerType], optional, defaults to "linear"): The learning rate scheduler type to use. See [SchedulerType] for all possible values. Common choices: - "linear" = [get_linear_schedule_with_warmup] - "cosine" = [get_cosine_schedule_with_warmup] - "constant" = [get_constant_schedule] - "constant_with_warmup" = [get_constant_schedule_with_warmup] lr_scheduler_kwargs (dict or str, optional, defaults to None): The extra arguments for the lr_scheduler. See the documentation of each scheduler for possible values. warmup_steps (int or float, optional, defaults to 0): Number of steps for a linear warmup from 0 to learning_rate. Warmup helps stabilize training in the initial phase. Can be: - An integer: exact number of warmup steps - A float in range [0, 1): interpreted as ratio of total training steps

Optimizer

optim (str or [training_args.OptimizerNames], optional, defaults to "adamw_torch" (for torch>=2.8 "adamw_torch_fused")): The optimizer to use. Common options: - "adamw_torch": PyTorch's AdamW (recommended default) - "adamw_torch_fused": Fused AdamW kernel - "adamw_hf": HuggingFace's AdamW implementation - "sgd": Stochastic Gradient Descent with momentum - "adafactor": Memory-efficient optimizer for large models - "adamw_8bit": 8-bit AdamW (requires bitsandbytes) See [OptimizerNames] for the complete list. optim_args (str, optional): Optional arguments that are supplied to optimizers such as AnyPrecisionAdamW, AdEMAMix, and GaLore. weight_decay (float, optional, defaults to 0): Weight decay coefficient applied by the optimizer (not the loss function). Adds L2 regularization to prevent overfitting by penalizing large weights. Automatically excluded from bias and LayerNorm parameters. Typical values: 0.01 (standard), 0.1 (stronger regularization), 0.0 (no regularization). adam_beta1 (float, optional, defaults to 0.9): The exponential decay rate for the first moment estimates (momentum) in Adam-based optimizers. Controls how much history of gradients to retain. adam_beta2 (float, optional, defaults to 0.999): The exponential decay rate for the second moment estimates (variance) in Adam-based optimizers. Controls adaptive learning rate scaling. adam_epsilon (float, optional, defaults to 1e-8): Epsilon value for numerical stability in Adam-based optimizers. Prevents division by zero in the denominator of the update rule. optim_target_modules (Union[str, list[str]], optional): The target modules to optimize, i.e. the module names that you would like to train. Currently used for the GaLore algorithm and APOLLO algorithm. See GaLore implementation and APOLLO implementation for more details. You need to make sure to pass a valid GaLore or APOLLO optimizer, e.g., one of: "apollo_adamw", "galore_adamw", "galore_adamw_8bit", "galore_adafactor" and make sure that the target modules are nn.Linear modules only.

Regularization & Training Stability

gradient_accumulation_steps (int, optional, defaults to 1): Number of update steps to accumulate gradients before performing a backward/update pass. Simulates larger batch sizes without additional memory. Effective batch size = per_device_train_batch_size × num_devices × gradient_accumulation_steps. > [!TIP] > When using gradient accumulation, one "step" is counted as one step with a backward pass. Therefore, logging, evaluation, and saving will occur every gradient_accumulation_steps × xxx_step training examples. average_tokens_across_devices (bool, optional, defaults to True): Whether or not to average tokens across devices. If enabled, will use all_reduce to synchronize num_tokens_in_batch for precise loss calculation. Reference: https://github.com/huggingface/transformers/issues/34242 max_grad_norm (float, optional, defaults to 1.0): Maximum gradient norm for gradient clipping. Applied after backward pass, before optimizer step. Prevents gradient explosion by scaling down gradients when their global norm exceeds this threshold. Set to 0 to disable clipping. Typical values: 1.0 (standard), 0.5 (more conservative), 5.0 (less aggressive). label_smoothing_factor (float, optional, defaults to 0.0): Label smoothing factor to prevent overconfidence. Replaces hard 0/1 targets with soft targets: 0 becomes ε/num_labels and 1 becomes 1 - ε + ε/num_labels, where ε = label_smoothing_factor. Zero means no smoothing. Typical range: 0.0 to 0.1.

Mixed Precision Training

bf16 (bool, optional, defaults to False): Enable bfloat16 (BF16) mixed precision training Generally preferred over FP16 due to better numerical stability and no loss scaling required. fp16 (bool, optional, defaults to False): Enable float16 (FP16) mixed precision training. Consider using BF16 instead if your hardware supports it. bf16_full_eval (bool, optional, defaults to False): Use full BF16 precision for evaluation (not just mixed precision). Faster and saves memory but may affect metric values slightly. Only applies during evaluation. fp16_full_eval (bool, optional, defaults to False): Use full FP16 precision for evaluation (not just mixed precision). Faster and saves memory but may affect metric values slightly. Only applies during evaluation. tf32 (bool, optional): Enable TensorFloat-32 (TF32) mode on Ampere and newer GPUs. TF32 uses 19-bit precision for matrix multiplications (instead of FP32's 23-bit), providing up to 8x speedup with negligible accuracy loss. Default depends on PyTorch version. See TF32 docs.

Gradient Checkpointing

gradient_checkpointing (bool, optional, defaults to False): Enable gradient checkpointing to trade compute for memory. Reduces memory usage by clearing activations during forward pass and recomputing them during backward pass. Enables training larger models or batch sizes at the cost of ~20% slower training. gradient_checkpointing_kwargs (dict, optional, defaults to None): Keyword arguments passed to gradient_checkpointing_enable().

Compilation

torch_compile (bool, optional, defaults to False): Compile the model using PyTorch 2.0's torch.compile() for faster training. Can provide 20-50% speedup with no code changes. Uses default compilation settings unless torch_compile_backend or torch_compile_mode are specified. torch_compile_backend (str, optional): Backend for torch.compile(). If set, automatically enables torch_compile. Options include "inductor" (default), "aot_eager", "cudagraphs". Backends vary by PyTorch version - see PyTorch docs for available options. torch_compile_mode (str, optional): Compilation mode for torch.compile(). If set, automatically enables torch_compile. Options: "default", "reduce-overhead" (minimize Python overhead), "max-autotune" (aggressive optimization, slower compile time).

Kernels

use_liger_kernel (bool, optional, defaults to False): Enable Liger Kernel optimizations. Increases multi-GPU throughput by ~20% and reduces memory usage by ~60%. Works with Flash Attention, FSDP, and DeepSpeed. Currently supports Llama, Mistral, Mixtral, and Gemma models. liger_kernel_config (Optional[dict], optional): Configuration for Liger Kernel. Passed as kwargs to _apply_liger_kernel_to_instance(). Options typically include: "rope", "swiglu", "cross_entropy", "fused_linear_cross_entropy", "rms_norm". If None, uses default configuration.

Additional Optimizations

use_cache (bool, optional, defaults to False): Whether or not to enable cache for the model. For training, this is usually not needed apart from some PEFT methods that uses past_key_values. neftune_noise_alpha (Optional[float]): If not None, this will activate NEFTune noise embeddings. This can drastically improve model performance for instruction fine-tuning. Check out the original paper and the original code. Support transformers PreTrainedModel and also PeftModel from peft. The original paper used values in the range [5.0, 15.0]. torch_empty_cache_steps (int, optional): Number of steps to wait before calling torch.<device>.empty_cache(). If left unset or set to None, cache will not be emptied. This can help avoid CUDA out-of-memory errors by lowering peak VRAM usage at a cost of about 10% slower performance. auto_find_batch_size (bool, optional, defaults to False) Whether to find a batch size that will fit into memory automatically through exponential decay, avoiding CUDA Out-of-Memory errors.

Logging & Monitoring Training

logging_strategy (str or [~trainer_utils.IntervalStrategy], optional, defaults to "steps"): The logging strategy to adopt during training. Possible values are: - "no": No logging is done during training. - "epoch": Logging is done at the end of each epoch. - "steps": Logging is done every logging_steps. logging_steps (int or float, optional, defaults to 500): Number of update steps between two logs if logging_strategy="steps". Should be an integer or a float in range [0,1). If smaller than 1, will be interpreted as ratio of total training steps. logging_first_step (bool, optional, defaults to False): Whether to log the first global_step or not. log_on_each_node (bool, optional, defaults to True): In multinode distributed training, whether to log using log_level once per node, or only on the main node. logging_nan_inf_filter (bool, optional, defaults to True): Filter out NaN and Inf losses when logging. If True, replaces NaN/Inf losses with the average of recent valid losses. Does not affect gradient computation, only logging. include_num_input_tokens_seen (Optional[Union[str, bool]], optional, defaults to "no"): Whether to track the number of input tokens seen. Must be one of ["all", "non_padding", "no"] or a boolean value which map to "all" or "no". May be slower in distributed training as gather operations must be called.

Logging

log_level (str, optional, defaults to passive): Logging level for the main process. Options: "debug", "info", "warning", "error", "critical", or "passive" (doesn't change the current Transformers logging level, which defaults to "warning") log_level_replica (str, optional, defaults to "warning"): Logging level for replica processes in distributed training. Same options as log_level. disable_tqdm (bool, optional): Disable tqdm progress bars. Defaults to True if log_level is warning or lower, False otherwise.

Experiment Tracking Integration

report_to (str or list[str], optional, defaults to "none"): The list of integrations to report the results and logs to. Supported platforms are "azure_ml", "clearml", "codecarbon", "comet_ml", "dagshub", "dvclive", "flyte", "mlflow", "swanlab", "tensorboard", "trackio" and "wandb". Use "all" to report to all integrations installed, "none" for no integrations. run_name (str, optional): A descriptor for the run. Typically used for trackio, wandb, mlflow, comet and swanlab logging. project (str, optional, defaults to "huggingface"): The name of the project to use for logging. Currently, only used by Trackio. trackio_space_id (str or None, optional, defaults to None): The Hugging Face Space ID to use for live Trackio logging with a Gradio-based Space. Should be a full Space name like 'username/reponame' or 'orgname/reponame', or just 'reponame' (the Space is created in the currently logged-in user's namespace). If None, metrics are logged only to a local directory (no Space on the Hub). That Gradio Space has read and write access to the Trackio Bucket, which is what you want while training is in progress—for example when resuming a partial run or aggregating logs from multiple machines. The Space will be public unless you set hub_private_repo=True or your organization's default is to create private Spaces. trackio_bucket_id (str or None, optional, defaults to None): Optional Hugging Face Bucket id for Trackio. If unset, Trackio derives one. Used together with a Gradio Space (trackio_space_id) and when deploying a static Space (trackio_static_space_id is not False). trackio_static_space_id (str, False, or None, optional, defaults to None): The Hugging Face Space ID to use for static Space created after training is complete. Should be a full Space name like 'username/reponame' or 'orgname/reponame', or just 'reponame' (the Space is created in the currently logged-in user's namespace). If False, no static Space will be created. If None, and model is pushed to the Hub, a static Space will be created with a default name and this will be linked from the model card. The Space will be public unless you set hub_private_repo=True or your organization's default is to create private Spaces.

Evaluation

eval_strategy (str or [~trainer_utils.IntervalStrategy], optional, defaults to "no"): When to run evaluation. Options: - "no": No evaluation during training - "steps": Evaluate every eval_steps - "epoch": Evaluate at the end of each epoch eval_steps (int or float, optional): Number of update steps between two evaluations if eval_strategy="steps". Will default to the same value as logging_steps if not set. Should be an integer or a float in range [0,1). If smaller than 1, will be interpreted as ratio of total training steps. eval_delay (float, optional): Number of epochs or steps to wait for before the first evaluation can be performed, depending on the eval_strategy. per_device_eval_batch_size (int, optional, defaults to 8): The batch size per device accelerator core/CPU for evaluation. prediction_loss_only (bool, optional, defaults to False): When performing evaluation and generating predictions, only returns the loss. eval_on_start (bool, optional, defaults to False): Whether to perform a evaluation step (sanity check) before the training to ensure the validation steps works correctly. eval_do_concat_batches (bool, optional, defaults to True): Whether to recursively concat inputs/losses/labels/predictions across batches. If False, will instead store them as lists, with each batch kept separate. eval_use_gather_object (bool, optional, defaults to False): Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices. This should only be enabled if users are not just returning tensors, and this is actively discouraged by PyTorch. This is useful when the labels structure is non standard, like in computer vision tasks. eval_accumulation_steps (int, optional): Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. If left unset, the whole predictions are accumulated on the device accelerator before being moved to the CPU (faster but requires more memory).

Metrics Computation

include_for_metrics (list[str], optional, defaults to []): Include additional data in the compute_metrics function if needed for metrics computation. Possible options to add to include_for_metrics list: - "inputs": Input data passed to the model, intended for calculating input dependent metrics. - "loss": Loss values computed during evaluation, intended for calculating loss dependent metrics. batch_eval_metrics (bool, optional, defaults to False): If set to True, evaluation will call compute_metrics at the end of each batch to accumulate statistics rather than saving all eval logits in memory. When set to True, you must pass a compute_metrics function that takes a boolean argument compute_result, which when passed True, will trigger the final global summary statistics from the batch-level summary statistics you've accumulated over the evaluation set.

Checkpointing & Saving

save_only_model (bool, optional, defaults to False): Save only model weights, not optimizer/scheduler/RNG state. Significantly reduces checkpoint size but prevents resuming training from the checkpoint. Use when you only need the trained model for inference, not continued training. You can only load the model using from_pretrained with this option set to True. save_strategy (str or [~trainer_utils.SaveStrategy], optional, defaults to "steps"): The checkpoint save strategy to adopt during training. Possible values are: - "no": No save is done during training. - "epoch": Save is done at the end of each epoch. - "steps": Save is done every save_steps. - "best": Save is done whenever a new best_metric is achieved. save_steps (int or float, optional, defaults to 500): Number of updates steps before two checkpoint saves if save_strategy="steps". Should be an integer or a float in range [0,1). If smaller than 1, will be interpreted as ratio of total training steps. save_on_each_node (bool, optional, defaults to False): When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on the main one. This should not be activated when the different nodes use the same storage as the files will be saved with the same names for each node. save_total_limit (int, optional): Maximum number of checkpoints to keep. Deletes older checkpoints in output_dir. When load_best_model_at_end=True, the best checkpoint is always retained plus the most recent ones. For example, save_total_limit=5 keeps the 4 most recent plus the best enable_jit_checkpoint (bool, optional, defaults to False): Enable Just-In-Time checkpointing on SIGTERM signal for graceful termination on preemptible workloads. Important: Configure your orchestrator's graceful shutdown period to allow sufficient time. For Kubernetes, set terminationGracePeriodSeconds (default 30s is usually insufficient). For Slurm, use --signal=USR1@<seconds>. Required grace period ≥ longest iteration time + checkpoint save time.

Hugging Face Hub Integration

push_to_hub (bool, optional, defaults to False): Whether or not to push the model to the Hub every time the model is saved. If this is activated, output_dir will begin a git directory synced with the repo (determined by hub_model_id) and the content will be pushed each time a save is triggered (depending on your save_strategy). Calling [~Trainer.save_model] will also trigger a push. hub_token (str, optional): The token to use to push the model to the Hub. Will default to the token in the cache folder obtained with hf auth login. hub_private_repo (bool, optional): Whether to make the repo private. If None (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. If reporting to Trackio with deployment to Hugging Face Spaces enabled, the same logic determines whether the Space is private. hub_model_id (str, optional): The name of the repository to keep in sync with the local output_dir. It can be a simple model ID in which case the model will be pushed in your namespace. Otherwise it should be the whole repository name, for instance "user_name/model", which allows you to push to an organization you are a member of with "organization_name/model". Will default to user_name/output_dir_name with output_dir_name being the name of output_dir. hub_strategy (str or [~trainer_utils.HubStrategy], optional, defaults to "every_save"): Defines what and when to push to Hub. Options: - "end": Push only at the end of training - "every_save": Push on each save (async to not block training) - "checkpoint": Like "every_save" plus push latest checkpoint to "last-checkpoint" subfolder for easy resuming - "all_checkpoints": Push all checkpoints as they appear hub_always_push (bool, optional, defaults to False): Unless this is True, the Trainer will skip pushing a checkpoint when the previous push is not finished. hub_revision (str, optional): The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash.

Best Model Tracking

load_best_model_at_end (bool, optional, defaults to False): Load the best checkpoint at the end of training. Requires eval_strategy to be set. When enabled, the best checkpoint is always saved (see save_total_limit). When True, save_strategy must match eval_strategy (unless save_strategy is "best"), and if using "steps", save_steps must be a multiple of eval_steps. metric_for_best_model (str, optional): Metric to use for comparing models when load_best_model_at_end=True. Must be a metric name returned by evaluation, with or without the "eval_" prefix. Defaults to "loss". If you set this, greater_is_better will default to True unless the name ends with "loss". Examples: "accuracy", "f1", "eval_bleu". greater_is_better (bool, optional): Whether higher metric values are better. Defaults based on metric_for_best_model: True if the metric name doesn't end in "loss", False otherwise.

Resuming Training

ignore_data_skip (bool, optional, defaults to False): When resuming training, skip fast-forwarding through the dataset to reach the previous state. If True, training starts from the beginning of the dataset (faster resume but results won't match interrupted training). If False, skips seen data (slower resume but exact continuation). restore_callback_states_from_checkpoint (bool, optional, defaults to False): Restore callback states from checkpoint when resuming. If True, will override callbacks passed to Trainer if they exist in the checkpoint.

Reproducibility

full_determinism (bool, optional, defaults to False) If True, [enable_full_determinism] is called instead of [set_seed] to ensure reproducible results in distributed training. Important: this will negatively impact the performance, so only use it for debugging. seed (int, optional, defaults to 42): Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use the [~Trainer.model_init] function to instantiate the model if it has some randomly initialized parameters. data_seed (int, optional): Random seed to be used with data samplers. If not set, random generators for data sampling will use the same seed as seed. This can be used to ensure reproducibility of data sampling, independent of the model seed.

Hardware Configuration

use_cpu (bool, optional, defaults to False): Whether or not to use cpu. If set to False, we will use the available torch device/backend.

Accelerate Configuration

accelerator_config (str, dict, or AcceleratorConfig, optional): Configuration for the internal Accelerate integration. Can be: - Path to JSON config file: "accelerator_config.json" - Dictionary with config options - AcceleratorConfig instance Key options: - split_batches (bool, defaults to False): Whether to split batches across devices. If True, actual batch size is the same on all devices (total must be divisible by num_processes). If False, each device gets the specified batch size. - dispatch_batches (bool): If True, only main process iterates through dataloader and dispatches batches to devices. Defaults to True for IterableDataset, False otherwise. - even_batches (bool, defaults to True): Duplicate samples from dataset start to ensure all workers get equal batch sizes. - use_seedable_sampler (bool, defaults to True): Use fully seedable random sampler for reproducibility. - use_configured_state (bool, defaults to False): Use pre-initialized AcceleratorState/PartialState instead of creating new one. May cause issues with hyperparameter tuning.

parallelism_config (ParallelismConfig, optional): Parallelism configuration for the training run. Requires Accelerate 1.10.1

Dataloader

dataloader_drop_last (bool, optional, defaults to False): Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size) or not. dataloader_num_workers (int, optional, defaults to 0): Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in the main process. dataloader_pin_memory (bool, optional, defaults to True): Whether you want to pin memory in data loaders or not. Will default to True. dataloader_persistent_workers (bool, optional, defaults to False): If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will increase RAM usage. Will default to False. dataloader_prefetch_factor (int, optional): Number of batches loaded in advance by each worker. 2 means there will be a total of 2 * num_workers batches prefetched across all workers. remove_unused_columns (bool, optional, defaults to True): Whether or not to automatically remove the columns unused by the model forward method. label_names (list[str], optional): The list of keys in your dictionary of inputs that correspond to the labels. Will eventually default to the list of argument names accepted by the model that contain the word "label", except if the model used is one of the XxxForQuestionAnswering in which case it will also include the ["start_positions", "end_positions"] keys. You should only specify label_names if you're using custom label names or if your model's forward consumes multiple label tensors (e.g., extractive QA). train_sampling_strategy (str, optional, defaults to "random"): The sampler to use for the training dataloader. Possible values are:

    - `"random"`: Uses <code>RandomSampler</code> (default).
    - `"sequential"`: Uses <code>SequentialSampler</code>.
    - `"group_by_length"`: Uses <code>LengthGroupedSampler</code> to group samples of roughly the same length
      together (to minimize padding and be more efficient).

Note: When using an <code>IterableDataset</code>, this argument is ignored.

length_column_name (str, optional, defaults to "length"): Column name for precomputed lengths. If the column exists, grouping by length will use these values rather than computing them on train startup. Ignored unless train_sampling_strategy is "group_by_length" and the dataset is an instance of Dataset.

DDP (DistributedDataParallel)

ddp_find_unused_parameters (bool, optional): When using distributed training, the value of the flag find_unused_parameters passed to DistributedDataParallel. Will default to False if gradient checkpointing is used, True otherwise. ddp_bucket_cap_mb (int, optional): When using distributed training, the value of the flag bucket_cap_mb passed to DistributedDataParallel. ddp_broadcast_buffers (bool, optional): When using distributed training, the value of the flag broadcast_buffers passed to DistributedDataParallel. Will default to False if gradient checkpointing is used, True otherwise. ddp_static_graph (bool, optional): When using distributed training, the value of the flag static_graph passed to DistributedDataParallel. ddp_backend (str, optional): The backend to use for distributed training. Must be one of "nccl", "mpi", "xccl", "gloo", "hccl". ddp_timeout (int, optional, defaults to 1800): The timeout for torch.distributed.init_process_group calls, used to avoid GPU socket timeouts when performing slow operations in distributed runnings. Please refer to the PyTorch documentation for more information.

FSDP (Fully Sharded Data Parallel)

fsdp (bool, str or list of [~trainer_utils.FSDPOption], optional, defaults to None): Enable PyTorch Fully Sharded Data Parallel (FSDP) for distributed training. Options: - "full_shard": Shard parameters, gradients, and optimizer states (most memory efficient) - "shard_grad_op": Shard only optimizer states and gradients (ZeRO-2) - "hybrid_shard": Full shard within nodes, replicate across nodes - "hybrid_shard_zero2": Shard gradients/optimizer within nodes, replicate across nodes - "offload": Offload parameters and gradients to CPU (only with "full_shard" or "shard_grad_op") - "auto_wrap": Automatically wrap layers using default_auto_wrap_policy fsdp_config (str or dict, optional): Config to be used with fsdp (Pytorch Distributed Parallel Training). The value is either a location of fsdp json config file (e.g., fsdp_config.json) or an already loaded json file as dict.

A List of config and its options:
    - fsdp_version (<code>int</code>, *optional*, defaults to <code>1</code>):
        The version of FSDP to use. Defaults to 1.
    - min_num_params (<code>int</code>, *optional*, defaults to <code>0</code>):
        FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when <code>fsdp</code> field is
        passed).
    - transformer_layer_cls_to_wrap (<code>list\[str]</code>, *optional*):
        List of transformer layer class names (case-sensitive) to wrap, e.g, <code>BertLayer</code>, <code>GPTJBlock</code>,
        <code>T5Block</code> .... (useful only when <code>fsdp</code> flag is passed).
    - backward_prefetch (<code>str</code>, *optional*)
        FSDP's backward prefetch mode. Controls when to prefetch next set of parameters (useful only when
        <code>fsdp</code> field is passed).

        A list of options along the following:

        - `"backward_pre"` : Prefetches the next set of parameters before the current set of parameter's
          gradient computation.
        - `"backward_post"` : This prefetches the next set of parameters after the current set of
          parameter's gradient computation.
    - forward_prefetch (<code>bool</code>, *optional*, defaults to <code>False</code>)
        FSDP's forward prefetch mode (useful only when <code>fsdp</code> field is passed).
         If `"True"`, then FSDP explicitly prefetches the next upcoming all-gather while executing in the
         forward pass.
    - limit_all_gathers (<code>bool</code>, *optional*, defaults to <code>False</code>)
        FSDP's limit_all_gathers (useful only when <code>fsdp</code> field is passed).
         If `"True"`, FSDP explicitly synchronizes the CPU thread to prevent too many in-flight
         all-gathers.
    - use_orig_params (<code>bool</code>, *optional*, defaults to <code>True</code>)
        If `"True"`, allows non-uniform <code>requires\_grad</code> during init, which means support for interspersed
        frozen and trainable parameters. Useful in cases such as parameter-efficient fine-tuning. Please
        refer this
        [blog](<https://dev-discuss.pytorch.org/t/rethinking-pytorch-fully-sharded-data-parallel-fsdp-from-first-principles/1019>
    - sync_module_states (<code>bool</code>, *optional*, defaults to <code>True</code>)
        If `"True"`, each individually wrapped FSDP unit will broadcast module parameters from rank 0 to
        ensure they are the same across all ranks after initialization
    - cpu_ram_efficient_loading (<code>bool</code>, *optional*, defaults to <code>False</code>)
        If `"True"`, only the first process loads the pretrained model checkpoint while all other processes
        have empty weights.  When this setting as `"True"`, <code>sync\_module\_states</code> also must to be `"True"`,
        otherwise all the processes except the main process would have random weights leading to unexpected
        behaviour during training.
    - activation_checkpointing (<code>bool</code>, *optional*, defaults to <code>False</code>):
        If `"True"`, activation checkpointing is a technique to reduce memory usage by clearing activations of
        certain layers and recomputing them during a backward pass. Effectively, this trades extra
        computation time for reduced memory usage.
    - xla (<code>bool</code>, *optional*, defaults to <code>False</code>):
        Whether to use PyTorch/XLA Fully Sharded Data Parallel Training. This is an experimental feature
        and its API may evolve in the future.
    - xla_fsdp_settings (<code>dict</code>, *optional*)
        The value is a dictionary which stores the XLA FSDP wrapping parameters.

        For a complete list of options, please see [here](
        <https://github.com/pytorch/xla/blob/master/torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py>).
    - xla_fsdp_grad_ckpt (<code>bool</code>, *optional*, defaults to <code>False</code>):
        Will use gradient checkpointing over each nested XLA FSDP wrapped layer. This setting can only be
        used when the xla flag is set to true, and an auto wrapping policy is specified through
        fsdp_min_num_params or fsdp_transformer_layer_cls_to_wrap.

DeepSpeed

deepspeed (str or dict, optional): Enable DeepSpeed integration. Value is either: - Path to DeepSpeed JSON config file: "ds_config.json" - Loaded config as dictionary

[!TIP] If using ZeRO initialization, instantiate your model after initializing TrainingArguments, otherwise ZeRO won't be applied.

Debugging & Profiling (Experimental)

debug (str or list of [~debug_utils.DebugOption], optional, defaults to ""): Enable one or more debug features. This is an experimental feature. Possible options are: - "underflow_overflow": detects overflow in model's input/outputs and reports the last frames that led to the event - "tpu_metrics_debug": print debug metrics on TPU skip_memory_metrics (bool, optional, defaults to True): Whether to skip adding of memory profiler reports to metrics. This is skipped by default because it slows down the training and evaluation speed.

External Script Flags (not used by Trainer)

do_train (bool, optional, defaults to False): Whether to run training or not. This argument is not directly used by [Trainer], it's intended to be used by your training/evaluation scripts instead. See the example scripts for more details. do_eval (bool, optional): Whether to run evaluation on the validation set or not. Will be set to True if eval_strategy is different from "no". This argument is not directly used by [Trainer], it's intended to be used by your training/evaluation scripts instead. See the example scripts for more details. do_predict (bool, optional, defaults to False): Whether to run predictions on the test set or not. This argument is not directly used by [Trainer], it's intended to be used by your training/evaluation scripts instead. See the example scripts for more details. resume_from_checkpoint (str, optional): The path to a folder with a valid checkpoint for your model. This argument is not directly used by [Trainer], it's intended to be used by your training/evaluation scripts instead. See the example scripts for more details.

Class variables

var default_optim

The type of the None singleton.

var is_torch_greater_or_equal_than_2_8

The type of the None singleton.

Instance variables

var accelerator_config : dict | str | None

The type of the None singleton.

var adam_beta1 : float

The type of the None singleton.

var adam_beta2 : float

The type of the None singleton.

var adam_epsilon : float

The type of the None singleton.

var auto_find_batch_size : bool

The type of the None singleton.

var average_tokens_across_devices : bool

The type of the None singleton.

var batch_eval_metrics : bool

The type of the None singleton.

var bf16 : bool

The type of the None singleton.

var bf16_full_eval : bool

The type of the None singleton.

var data_seed : int | None

The type of the None singleton.

var dataloader_drop_last : bool

The type of the None singleton.

var dataloader_num_workers : int

The type of the None singleton.

var dataloader_persistent_workers : bool

The type of the None singleton.

var dataloader_pin_memory : bool

The type of the None singleton.

var dataloader_prefetch_factor : int | None

The type of the None singleton.

var ddp_backend : str | None

The type of the None singleton.

var ddp_broadcast_buffers : bool | None

The type of the None singleton.

var ddp_bucket_cap_mb : int | None

The type of the None singleton.

var ddp_find_unused_parameters : bool | None

The type of the None singleton.

var ddp_static_graph : bool | None

The type of the None singleton.

var ddp_timeout : int

The type of the None singleton.

prop ddp_timeout_delta : datetime.timedelta

Expand source code

@property
def ddp_timeout_delta(self) -> timedelta:
    """
    The actual timeout for torch.distributed.init_process_group since it expects a timedelta variable.
    """
    return timedelta(seconds=self.ddp_timeout)

The actual timeout for torch.distributed.init_process_group since it expects a timedelta variable.

var debug : str | list[transformers.debug_utils.DebugOption]

The type of the None singleton.

var deepspeed : dict | str | None

The type of the None singleton.

prop device : torch.device

Expand source code

@property
def device(self) -> "torch.device":
    """
    The device used by this process.
    """
    requires_backends(self, ["torch"])
    return self._setup_devices

The device used by this process.

var disable_tqdm : bool | None

The type of the None singleton.

var do_eval : bool

The type of the None singleton.

var do_predict : bool

The type of the None singleton.

var do_train : bool

The type of the None singleton.

var enable_jit_checkpoint : bool

The type of the None singleton.

var eval_accumulation_steps : int | None

The type of the None singleton.

prop eval_batch_size : int

Expand source code

@property
def eval_batch_size(self) -> int:
    """
    The actual batch size for evaluation.
    """
    eval_batch_size = self.per_device_eval_batch_size * max(1, self.n_gpu)
    return eval_batch_size

The actual batch size for evaluation.

var eval_delay : float

The type of the None singleton.

var eval_do_concat_batches : bool

The type of the None singleton.

var eval_on_start : bool

The type of the None singleton.

var eval_steps : float | None

The type of the None singleton.

var eval_strategy : transformers.trainer_utils.IntervalStrategy | str

The type of the None singleton.

var eval_use_gather_object : bool

The type of the None singleton.

var fp16 : bool

The type of the None singleton.

var fp16_full_eval : bool

The type of the None singleton.

var fsdp : list[transformers.trainer_utils.FSDPOption] | str | None

The type of the None singleton.

var fsdp_config : dict[str, typing.Any] | str | None

The type of the None singleton.

var full_determinism : bool

The type of the None singleton.

var gradient_accumulation_steps : int

The type of the None singleton.

var gradient_checkpointing : bool

The type of the None singleton.

var gradient_checkpointing_kwargs : dict[str, typing.Any] | str | None

The type of the None singleton.

var greater_is_better : bool | None

The type of the None singleton.

var hub_always_push : bool

The type of the None singleton.

var hub_model_id : str | None

The type of the None singleton.

var hub_private_repo : bool | None

The type of the None singleton.

var hub_revision : str | None

The type of the None singleton.

var hub_strategy : transformers.trainer_utils.HubStrategy | str

The type of the None singleton.

var hub_token : str | None

The type of the None singleton.

var ignore_data_skip : bool

The type of the None singleton.

var include_for_metrics : list[str]

The type of the None singleton.

var include_num_input_tokens_seen : str | bool

The type of the None singleton.

var label_names : list[str] | None

The type of the None singleton.

var label_smoothing_factor : float

The type of the None singleton.

var learning_rate : float

The type of the None singleton.

var length_column_name : str

The type of the None singleton.

var liger_kernel_config : dict[str, bool] | None

The type of the None singleton.

var load_best_model_at_end : bool

The type of the None singleton.

prop local_process_index

Expand source code

@property
def local_process_index(self):
    """
    The index of the local process used.
    """
    requires_backends(self, ["torch"])

    if self.distributed_state is not None:
        return self.distributed_state.local_process_index
    elif is_sagemaker_mp_enabled():
        return smp.local_rank()
    return 0

The index of the local process used.

var local_rank : int

The type of the None singleton.

var log_level : str

The type of the None singleton.

var log_level_replica : str

The type of the None singleton.

var log_on_each_node : bool

The type of the None singleton.

var logging_dir : str | None

The type of the None singleton.

var logging_first_step : bool

The type of the None singleton.

var logging_nan_inf_filter : bool

The type of the None singleton.

var logging_steps : float

The type of the None singleton.

var logging_strategy : transformers.trainer_utils.IntervalStrategy | str

The type of the None singleton.

var lr_scheduler_kwargs : dict | str | None

The type of the None singleton.

var lr_scheduler_type : transformers.trainer_utils.SchedulerType | str

The type of the None singleton.

var max_grad_norm : float

The type of the None singleton.

var max_steps : int

The type of the None singleton.

var metric_for_best_model : str | None

The type of the None singleton.

prop n_gpu

Expand source code

@property
def n_gpu(self):
    """
    The number of GPUs used by this process.

    Note:
        This will only be greater than one when you have multiple GPUs available but are not using distributed
        training. For distributed training, it will always be 1.
    """
    requires_backends(self, ["torch"])
    # Ensure _setup_devices has been called
    if not hasattr(self, "_n_gpu"):
        _ = self._setup_devices
    return self._n_gpu

The number of GPUs used by this process.

Note -----= This will only be greater than one when you have multiple GPUs available but are not using distributed training. For distributed training, it will always be 1.

var neftune_noise_alpha : float | None

The type of the None singleton.

var num_train_epochs : float

The type of the None singleton.

var optim : transformers.training_args.OptimizerNames | str

The type of the None singleton.

var optim_args : str | None

The type of the None singleton.

var optim_target_modules : str | list[str] | None

The type of the None singleton.

var output_dir : str | None

The type of the None singleton.

prop parallel_mode

Expand source code

@property
def parallel_mode(self):
    """
    The current mode used for parallelism if multiple GPUs/TPU cores are available. One of:

    - `ParallelMode.NOT_PARALLEL`: no parallelism (CPU or one GPU).
    - `ParallelMode.NOT_DISTRIBUTED`: several GPUs in one single process (uses `torch.nn.DataParallel`).
    - `ParallelMode.DISTRIBUTED`: several GPUs, each having its own process (uses
      `torch.nn.DistributedDataParallel`).
    - `ParallelMode.TPU`: several TPU cores.
    """
    requires_backends(self, ["torch"])
    if is_torch_xla_available():
        return ParallelMode.TPU
    elif is_sagemaker_mp_enabled():
        return ParallelMode.SAGEMAKER_MODEL_PARALLEL
    elif is_sagemaker_dp_enabled():
        return ParallelMode.SAGEMAKER_DATA_PARALLEL
    elif self.distributed_state is not None and self.distributed_state.distributed_type != DistributedType.NO:
        return ParallelMode.DISTRIBUTED
    elif self.n_gpu > 1:
        return ParallelMode.NOT_DISTRIBUTED
    else:
        return ParallelMode.NOT_PARALLEL

The current mode used for parallelism if multiple GPUs/TPU cores are available. One of:

• ParallelMode.NOT_PARALLEL: no parallelism (CPU or one GPU).
• ParallelMode.NOT_DISTRIBUTED: several GPUs in one single process (uses torch.nn.DataParallel).
• ParallelMode.DISTRIBUTED: several GPUs, each having its own process (uses torch.nn.DistributedDataParallel).
• ParallelMode.TPU: several TPU cores.

var parallelism_config : accelerate.parallelism_config.ParallelismConfig | None

The type of the None singleton.

var per_device_eval_batch_size : int

The type of the None singleton.

var per_device_train_batch_size : int

The type of the None singleton.

prop place_model_on_device : bool | None

Expand source code

@property
def place_model_on_device(self) -> bool | None:
    """
    Can be subclassed and overridden for some specific integrations.
    """
    return None

Can be subclassed and overridden for some specific integrations.

var prediction_loss_only : bool

The type of the None singleton.

prop process_index

Expand source code

@property
def process_index(self):
    """
    The index of the current process used.
    """
    requires_backends(self, ["torch"])
    if self.distributed_state is not None:
        return self.distributed_state.process_index
    elif is_sagemaker_mp_enabled():
        return smp.dp_rank() if not smp.state.cfg.prescaled_batch else smp.rdp_rank()
    return 0

The index of the current process used.

var project : str

The type of the None singleton.

var push_to_hub : bool

The type of the None singleton.

var remove_unused_columns : bool

The type of the None singleton.

var report_to : str | list[str] | None

The type of the None singleton.

var restore_callback_states_from_checkpoint : bool

The type of the None singleton.

var resume_from_checkpoint : str | None

The type of the None singleton.

var run_name : str | None

The type of the None singleton.

var save_on_each_node : bool

The type of the None singleton.

var save_only_model : bool

The type of the None singleton.

var save_steps : float

The type of the None singleton.

var save_strategy : transformers.trainer_utils.SaveStrategy | str

The type of the None singleton.

var save_total_limit : int | None

The type of the None singleton.

var seed : int

The type of the None singleton.

prop should_log

Expand source code

@property
def should_log(self):
    """
    Whether or not the current process should produce log.
    """
    if self.log_on_each_node:
        return self.local_process_index == 0
    else:
        if is_sagemaker_mp_enabled():
            return smp.rank() == 0
        else:
            return self.process_index == 0

Whether or not the current process should produce log.

prop should_save

Expand source code

@property
def should_save(self):
    """
    Whether or not the current process should write to disk, e.g., to save models and checkpoints.
    """
    if self.save_on_each_node:
        return self.local_process_index == 0
    else:
        if is_sagemaker_mp_enabled():
            return smp.rank() == 0
        else:
            return self.process_index == 0

Whether or not the current process should write to disk, e.g., to save models and checkpoints.

var skip_memory_metrics : bool

The type of the None singleton.

var tf32 : bool | None

The type of the None singleton.

var torch_compile : bool

The type of the None singleton.

var torch_compile_backend : str | None

The type of the None singleton.

var torch_compile_mode : str | None

The type of the None singleton.

var torch_empty_cache_steps : int | None

The type of the None singleton.

var trackio_bucket_id : str | None

The type of the None singleton.

var trackio_space_id : str | None

The type of the None singleton.

var trackio_static_space_id : str | Literal[False] | None

The type of the None singleton.

prop train_batch_size : int

Expand source code

@property
def train_batch_size(self) -> int:
    """
    The actual batch size for training.
    """
    train_batch_size = self.per_device_train_batch_size * max(1, self.n_gpu)
    return train_batch_size

The actual batch size for training.

var train_sampling_strategy : str

The type of the None singleton.

var use_cache : bool

The type of the None singleton.

var use_cpu : bool

The type of the None singleton.

var use_liger_kernel : bool

The type of the None singleton.

var warmup_ratio : float | None

The type of the None singleton.

var warmup_steps : float

The type of the None singleton.

var weight_decay : float

The type of the None singleton.

prop world_size

Expand source code

@property
def world_size(self):
    """
    The number of processes used in parallel.
    """
    requires_backends(self, ["torch"])
    if self.distributed_state is not None:
        return self.distributed_state.num_processes
    elif is_sagemaker_mp_enabled():
        return smp.dp_size() if not smp.state.cfg.prescaled_batch else smp.rdp_size()
    return 1

The number of processes used in parallel.

Methods

def get_process_log_level(self)

Expand source code

def get_process_log_level(self):
    """
    Returns the log level to be used depending on whether this process is the main process of node 0, main process
    of node non-0, or a non-main process.

    For the main process the log level defaults to the logging level set (`logging.WARNING` if you didn't do
    anything) unless overridden by `log_level` argument.

    For the replica processes the log level defaults to `logging.WARNING` unless overridden by `log_level_replica`
    argument.

    The choice between the main and replica process settings is made according to the return value of `should_log`.
    """

    # convert to int
    log_level = trainer_log_levels[self.log_level]
    log_level_replica = trainer_log_levels[self.log_level_replica]

    log_level_main_node = logging.get_verbosity() if log_level == -1 else log_level
    log_level_replica_node = logging.get_verbosity() if log_level_replica == -1 else log_level_replica
    return log_level_main_node if self.should_log else log_level_replica_node

Returns the log level to be used depending on whether this process is the main process of node 0, main process of node non-0, or a non-main process.

For the main process the log level defaults to the logging level set (logging.WARNING if you didn't do anything) unless overridden by log_level argument.

For the replica processes the log level defaults to logging.WARNING unless overridden by log_level_replica argument.

The choice between the main and replica process settings is made according to the return value of should_log.

def get_warmup_steps(self, num_training_steps: int)

Expand source code

def get_warmup_steps(self, num_training_steps: int):
    """
    Get number of steps used for a linear warmup.
    """
    warmup_steps = (
        int(self.warmup_steps) if self.warmup_steps >= 1 else math.ceil(num_training_steps * self.warmup_steps)
    )
    return warmup_steps

Get number of steps used for a linear warmup.

def main_process_first(self, local=True, desc='work')

Expand source code

@contextlib.contextmanager
def main_process_first(self, local=True, desc="work"):
    """
    A context manager for torch distributed environment where on needs to do something on the main process, while
    blocking replicas, and when it's finished releasing the replicas.

    One such use is for `datasets`'s `map` feature which to be efficient should be run once on the main process,
    which upon completion saves a cached version of results and which then automatically gets loaded by the
    replicas.

    Args:
        local (`bool`, *optional*, defaults to `True`):
            if `True` first means process of rank 0 of each node if `False` first means process of rank 0 of node
            rank 0 In multi-node environment with a shared filesystem you most likely will want to use
            `local=False` so that only the main process of the first node will do the processing. If however, the
            filesystem is not shared, then the main process of each node will need to do the processing, which is
            the default behavior.
        desc (`str`, *optional*, defaults to `"work"`):
            a work description to be used in debug logs

    """
    if is_torch_available() and self.world_size > 1:
        main_process_desc = "main local process" if local else "main process"
        if self.distributed_state is not None:
            is_main_process = (
                self.distributed_state.is_local_main_process if local else self.distributed_state.is_main_process
            )
        elif is_sagemaker_mp_enabled():
            is_main_process = smp.rank() == 0

        try:
            if not is_main_process:
                # tell all replicas to wait
                logger.debug(f"{self.process_index}: waiting for the {main_process_desc} to perform {desc}")

                if is_torch_xla_available():
                    xm.rendezvous(desc)
                else:
                    dist.barrier()
            yield
        finally:
            if is_main_process:
                # the wait is over
                logger.debug(f"{self.process_index}: {main_process_desc} completed {desc}, releasing all replicas")
                if is_torch_xla_available():
                    xm.rendezvous(desc)
                else:
                    dist.barrier()
    else:
        yield

A context manager for torch distributed environment where on needs to do something on the main process, while blocking replicas, and when it's finished releasing the replicas.

One such use is for datasets's map feature which to be efficient should be run once on the main process, which upon completion saves a cached version of results and which then automatically gets loaded by the replicas.

Args -----= local (bool, optional, defaults to True): if True first means process of rank 0 of each node if False first means process of rank 0 of node rank 0 In multi-node environment with a shared filesystem you most likely will want to use local=False so that only the main process of the first node will do the processing. If however, the filesystem is not shared, then the main process of each node will need to do the processing, which is the default behavior. desc (str, optional, defaults to "work"): a work description to be used in debug logs

def set_dataloader(self,
train_batch_size: int = 8,
eval_batch_size: int = 8,
drop_last: bool = False,
num_workers: int = 0,
pin_memory: bool = True,
persistent_workers: bool = False,
prefetch_factor: int | None = None,
auto_find_batch_size: bool = False,
ignore_data_skip: bool = False,
sampler_seed: int | None = None)

Expand source code

def set_dataloader(
    self,
    train_batch_size: int = 8,
    eval_batch_size: int = 8,
    drop_last: bool = False,
    num_workers: int = 0,
    pin_memory: bool = True,
    persistent_workers: bool = False,
    prefetch_factor: int | None = None,
    auto_find_batch_size: bool = False,
    ignore_data_skip: bool = False,
    sampler_seed: int | None = None,
):
    """
    A method that regroups all arguments linked to the dataloaders creation.

    Args:
        drop_last (`bool`, *optional*, defaults to `False`):
            Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch
            size) or not.
        num_workers (`int`, *optional*, defaults to 0):
            Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in
            the main process.
        pin_memory (`bool`, *optional*, defaults to `True`):
            Whether you want to pin memory in data loaders or not. Will default to `True`.
        persistent_workers (`bool`, *optional*, defaults to `False`):
            If True, the data loader will not shut down the worker processes after a dataset has been consumed
            once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training,
            but will increase RAM usage. Will default to `False`.
        prefetch_factor (`int`, *optional*):
            Number of batches loaded in advance by each worker.
            2 means there will be a total of 2 * num_workers batches prefetched across all workers.
        auto_find_batch_size (`bool`, *optional*, defaults to `False`)
            Whether to find a batch size that will fit into memory automatically through exponential decay,
            avoiding CUDA Out-of-Memory errors. Requires accelerate to be installed (`pip install accelerate`)
        ignore_data_skip (`bool`, *optional*, defaults to `False`):
            When resuming training, whether or not to skip the epochs and batches to get the data loading at the
            same stage as in the previous training. If set to `True`, the training will begin faster (as that
            skipping step can take a long time) but will not yield the same results as the interrupted training
            would have.
        sampler_seed (`int`, *optional*):
            Random seed to be used with data samplers. If not set, random generators for data sampling will use the
            same seed as `self.seed`. This can be used to ensure reproducibility of data sampling, independent of
            the model seed.

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_dataloader(train_batch_size=16, eval_batch_size=64)
    >>> args.per_device_train_batch_size
    16
    ```
    """
    self.per_device_train_batch_size = train_batch_size
    self.per_device_eval_batch_size = eval_batch_size
    self.dataloader_drop_last = drop_last
    self.dataloader_num_workers = num_workers
    self.dataloader_pin_memory = pin_memory
    self.dataloader_persistent_workers = persistent_workers
    self.dataloader_prefetch_factor = prefetch_factor
    self.auto_find_batch_size = auto_find_batch_size
    self.ignore_data_skip = ignore_data_skip
    self.data_seed = sampler_seed
    return self

A method that regroups all arguments linked to the dataloaders creation.

Args -----= drop_last (bool, optional, defaults to False): Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size) or not. num_workers (int, optional, defaults to 0): Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in the main process. pin_memory (bool, optional, defaults to True): Whether you want to pin memory in data loaders or not. Will default to True. persistent_workers (bool, optional, defaults to False): If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will increase RAM usage. Will default to False. prefetch_factor (int, optional): Number of batches loaded in advance by each worker. 2 means there will be a total of 2 * num_workers batches prefetched across all workers. auto_find_batch_size (bool, optional, defaults to False) Whether to find a batch size that will fit into memory automatically through exponential decay, avoiding CUDA Out-of-Memory errors. Requires accelerate to be installed (pip install accelerate) ignore_data_skip (bool, optional, defaults to False): When resuming training, whether or not to skip the epochs and batches to get the data loading at the same stage as in the previous training. If set to True, the training will begin faster (as that skipping step can take a long time) but will not yield the same results as the interrupted training would have. sampler_seed (int, optional): Random seed to be used with data samplers. If not set, random generators for data sampling will use the same seed as self.seed. This can be used to ensure reproducibility of data sampling, independent of the model seed. Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_dataloader(train_batch_size=16, eval_batch_size=64)
>>> args.per_device_train_batch_size
16

def set_evaluate(self,
strategy: transformers.trainer_utils.IntervalStrategy | str = 'no',
steps: int = 500,
batch_size: int = 8,
accumulation_steps: int | None = None,
delay: float | None = None,
loss_only: bool = False)

Expand source code

def set_evaluate(
    self,
    strategy: str | IntervalStrategy = "no",
    steps: int = 500,
    batch_size: int = 8,
    accumulation_steps: int | None = None,
    delay: float | None = None,
    loss_only: bool = False,
):
    """
    A method that regroups all arguments linked to evaluation.

    Args:
        strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"no"`):
            The evaluation strategy to adopt during training. Possible values are:

                - `"no"`: No evaluation is done during training.
                - `"steps"`: Evaluation is done (and logged) every `steps`.
                - `"epoch"`: Evaluation is done at the end of each epoch.

            Setting a `strategy` different from `"no"` will set `self.do_eval` to `True`.
        steps (`int`, *optional*, defaults to 500):
            Number of update steps between two evaluations if `strategy="steps"`.
        batch_size (`int` *optional*, defaults to 8):
            The batch size per device (GPU/TPU core/CPU...) used for evaluation.
        accumulation_steps (`int`, *optional*):
            Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU.
            If left unset, the whole predictions are accumulated on GPU/TPU before being moved to the CPU (faster
            but requires more memory).
        delay (`float`, *optional*):
            Number of epochs or steps to wait for before the first evaluation can be performed, depending on the
            eval_strategy.
        loss_only (`bool`, *optional*, defaults to `False`):
            Ignores all outputs except the loss.

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_evaluate(strategy="steps", steps=100)
    >>> args.eval_steps
    100
    ```
    """
    self.eval_strategy = IntervalStrategy(strategy)
    if self.eval_strategy == IntervalStrategy.STEPS and steps == 0:
        raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.")
    self.do_eval = self.eval_strategy != IntervalStrategy.NO
    self.eval_steps = steps
    self.per_device_eval_batch_size = batch_size
    self.eval_accumulation_steps = accumulation_steps
    self.eval_delay = delay
    self.prediction_loss_only = loss_only
    return self

A method that regroups all arguments linked to evaluation.

Args -----= strategy (str or [~trainer_utils.IntervalStrategy], optional, defaults to "no"): The evaluation strategy to adopt during training. Possible values are:

    - `"no"`: No evaluation is done during training.
    - `"steps"`: Evaluation is done (and logged) every <code>steps</code>.
    - `"epoch"`: Evaluation is done at the end of each epoch.

Setting a <code>strategy</code> different from `"no"` will set <code>self.do\_eval</code> to <code>True</code>.

steps (int, optional, defaults to 500): Number of update steps between two evaluations if strategy="steps". batch_size (int optional, defaults to 8): The batch size per device (GPU/TPU core/CPU…) used for evaluation. accumulation_steps (int, optional): Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. If left unset, the whole predictions are accumulated on GPU/TPU before being moved to the CPU (faster but requires more memory). delay (float, optional): Number of epochs or steps to wait for before the first evaluation can be performed, depending on the eval_strategy. loss_only (bool, optional, defaults to False): Ignores all outputs except the loss. Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_evaluate(strategy="steps", steps=100)
>>> args.eval_steps
100

def set_logging(self,
strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps',
steps: int = 500,
report_to: str | list[str] = 'none',
level: str = 'passive',
first_step: bool = False,
nan_inf_filter: bool = False,
on_each_node: bool = False,
replica_level: str = 'passive')

Expand source code

def set_logging(
    self,
    strategy: str | IntervalStrategy = "steps",
    steps: int = 500,
    report_to: str | list[str] = "none",
    level: str = "passive",
    first_step: bool = False,
    nan_inf_filter: bool = False,
    on_each_node: bool = False,
    replica_level: str = "passive",
):
    """
    A method that regroups all arguments linked to logging.

    Args:
        strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`):
            The logging strategy to adopt during training. Possible values are:

                - `"no"`: No logging is done during training.
                - `"epoch"`: Logging is done at the end of each epoch.
                - `"steps"`: Logging is done every `logging_steps`.

        steps (`int`, *optional*, defaults to 500):
            Number of update steps between two logs if `strategy="steps"`.
        level (`str`, *optional*, defaults to `"passive"`):
            Logger log level to use on the main process. Possible choices are the log levels as strings: `"debug"`,
            `"info"`, `"warning"`, `"error"` and `"critical"`, plus a `"passive"` level which doesn't set anything
            and lets the application set the level.
        report_to (`str` or `list[str]`, *optional*, defaults to `"none"`):
            The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`,
            `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`,
            `"swanlab"`, `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations
            installed, `"none"` for no integrations.
        first_step (`bool`, *optional*, defaults to `False`):
            Whether to log and evaluate the first `global_step` or not.
        nan_inf_filter (`bool`, *optional*, defaults to `True`):
            Whether to filter `nan` and `inf` losses for logging. If set to `True` the loss of every step that is
            `nan` or `inf` is filtered and the average loss of the current logging window is taken instead.

            <Tip>

            `nan_inf_filter` only influences the logging of loss values, it does not change the behavior the
            gradient is computed or applied to the model.

            </Tip>

        on_each_node (`bool`, *optional*, defaults to `True`):
            In multinode distributed training, whether to log using `log_level` once per node, or only on the main
            node.
        replica_level (`str`, *optional*, defaults to `"passive"`):
            Logger log level to use on replicas. Same choices as `log_level`

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_logging(strategy="steps", steps=100)
    >>> args.logging_steps
    100
    ```
    """
    self.logging_strategy = IntervalStrategy(strategy)
    if self.logging_strategy == IntervalStrategy.STEPS and steps == 0:
        raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.")
    self.logging_steps = steps
    self.report_to = report_to
    self.log_level = level
    self.logging_first_step = first_step
    self.logging_nan_inf_filter = nan_inf_filter
    self.log_on_each_node = on_each_node
    self.log_level_replica = replica_level
    return self

A method that regroups all arguments linked to logging.

Args -----= strategy (str or [~trainer_utils.IntervalStrategy], optional, defaults to "steps"): The logging strategy to adopt during training. Possible values are:

    - `"no"`: No logging is done during training.
    - `"epoch"`: Logging is done at the end of each epoch.
    - `"steps"`: Logging is done every <code>logging\_steps</code>.

steps (int, optional, defaults to 500): Number of update steps between two logs if strategy="steps". level (str, optional, defaults to "passive"): Logger log level to use on the main process. Possible choices are the log levels as strings: "debug", "info", "warning", "error" and "critical", plus a "passive" level which doesn't set anything and lets the application set the level. report_to (str or list[str], optional, defaults to "none"): The list of integrations to report the results and logs to. Supported platforms are "azure_ml", "clearml", "codecarbon", "comet_ml", "dagshub", "dvclive", "flyte", "mlflow", "swanlab", "tensorboard", "trackio" and "wandb". Use "all" to report to all integrations installed, "none" for no integrations. first_step (bool, optional, defaults to False): Whether to log and evaluate the first global_step or not. nan_inf_filter (bool, optional, defaults to True): Whether to filter nan and inf losses for logging. If set to True the loss of every step that is nan or inf is filtered and the average loss of the current logging window is taken instead.

<Tip>

<code>nan\_inf\_filter</code> only influences the logging of loss values, it does not change the behavior the
gradient is computed or applied to the model.

</Tip>

on_each_node (bool, optional, defaults to True): In multinode distributed training, whether to log using log_level once per node, or only on the main node. replica_level (str, optional, defaults to "passive"): Logger log level to use on replicas. Same choices as log_level Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_logging(strategy="steps", steps=100)
>>> args.logging_steps
100

def set_lr_scheduler(self,
name: transformers.trainer_utils.SchedulerType | str = 'linear',
num_epochs: float = 3.0,
max_steps: int = -1,
warmup_steps: float = 0,
warmup_ratio: float | None = None)

Expand source code

def set_lr_scheduler(
    self,
    name: str | SchedulerType = "linear",
    num_epochs: float = 3.0,
    max_steps: int = -1,
    warmup_steps: float = 0,
    warmup_ratio: float | None = None,
):
    """
    A method that regroups all arguments linked to the learning rate scheduler and its hyperparameters.

    Args:
        name (`str` or [`SchedulerType`], *optional*, defaults to `"linear"`):
            The scheduler type to use. See the documentation of [`SchedulerType`] for all possible values.
        num_epochs(`float`, *optional*, defaults to 3.0):
            Total number of training epochs to perform (if not an integer, will perform the decimal part percents
            of the last epoch before stopping training).
        max_steps (`int`, *optional*, defaults to -1):
            If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`.
            For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until
            `max_steps` is reached.
        warmup_steps (`float`, *optional*, defaults to 0):
            Number of steps used for a linear warmup from 0 to `learning_rate`.  Should be an integer or a float in range `[0,1)`.
            If smaller than 1, will be interpreted as ratio of steps used for a linear warmup from 0 to `learning_rate`.

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_lr_scheduler(name="cosine", warmup_steps=0.05)
    >>> args.warmup_steps
    0.05
    ```
    """
    if warmup_ratio is not None:
        logger.warning("warmup_ratio is deprecated and will be removed in v5.2 . Use `warmup_steps` instead.")
        warmup_steps = warmup_ratio

    self.lr_scheduler_type = SchedulerType(name)
    self.num_train_epochs = num_epochs
    self.max_steps = max_steps
    self.warmup_steps = warmup_steps
    return self

A method that regroups all arguments linked to the learning rate scheduler and its hyperparameters.

Args -----= name (str or [SchedulerType], optional, defaults to "linear"): The scheduler type to use. See the documentation of [SchedulerType] for all possible values. num_epochs(float, optional, defaults to 3.0): Total number of training epochs to perform (if not an integer, will perform the decimal part percents of the last epoch before stopping training). max_steps (int, optional, defaults to -1): If set to a positive number, the total number of training steps to perform. Overrides num_train_epochs. For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until max_steps is reached. warmup_steps (float, optional, defaults to 0): Number of steps used for a linear warmup from 0 to learning_rate. Should be an integer or a float in range [0,1). If smaller than 1, will be interpreted as ratio of steps used for a linear warmup from 0 to learning_rate. Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_lr_scheduler(name="cosine", warmup_steps=0.05)
>>> args.warmup_steps
0.05

def set_optimizer(self,
name: transformers.training_args.OptimizerNames | str = 'adamw_torch',
learning_rate: float = 5e-05,
weight_decay: float = 0,
beta1: float = 0.9,
beta2: float = 0.999,
epsilon: float = 1e-08,
args: str | None = None)

Expand source code

def set_optimizer(
    self,
    name: str | OptimizerNames = "adamw_torch",
    learning_rate: float = 5e-5,
    weight_decay: float = 0,
    beta1: float = 0.9,
    beta2: float = 0.999,
    epsilon: float = 1e-8,
    args: str | None = None,
):
    """
    A method that regroups all arguments linked to the optimizer and its hyperparameters.

    Args:
        name (`str` or [`training_args.OptimizerNames`], *optional*, defaults to `"adamw_torch"`):
            The optimizer to use: `"adamw_torch"`, `"adamw_torch_fused"`, `"adamw_apex_fused"`,
            `"adamw_anyprecision"` or `"adafactor"`.
        learning_rate (`float`, *optional*, defaults to 5e-5):
            The initial learning rate.
        weight_decay (`float`, *optional*, defaults to 0):
            The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights.
        beta1 (`float`, *optional*, defaults to 0.9):
            The beta1 hyperparameter for the adam optimizer or its variants.
        beta2 (`float`, *optional*, defaults to 0.999):
            The beta2 hyperparameter for the adam optimizer or its variants.
        epsilon (`float`, *optional*, defaults to 1e-8):
            The epsilon hyperparameter for the adam optimizer or its variants.
        args (`str`, *optional*):
            Optional arguments that are supplied to AnyPrecisionAdamW (only useful when
            `optim="adamw_anyprecision"`).

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_optimizer(name="adamw_torch", beta1=0.8)
    >>> args.optim
    'adamw_torch'
    ```
    """
    self.optim = OptimizerNames(name)
    self.learning_rate = learning_rate
    self.weight_decay = weight_decay
    self.adam_beta1 = beta1
    self.adam_beta2 = beta2
    self.adam_epsilon = epsilon
    self.optim_args = args
    return self

A method that regroups all arguments linked to the optimizer and its hyperparameters.

Args -----= name (str or [training_args.OptimizerNames], optional, defaults to "adamw_torch"): The optimizer to use: "adamw_torch", "adamw_torch_fused", "adamw_apex_fused", "adamw_anyprecision" or "adafactor". learning_rate (float, optional, defaults to 5e-5): The initial learning rate. weight_decay (float, optional, defaults to 0): The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights. beta1 (float, optional, defaults to 0.9): The beta1 hyperparameter for the adam optimizer or its variants. beta2 (float, optional, defaults to 0.999): The beta2 hyperparameter for the adam optimizer or its variants. epsilon (float, optional, defaults to 1e-8): The epsilon hyperparameter for the adam optimizer or its variants. args (str, optional): Optional arguments that are supplied to AnyPrecisionAdamW (only useful when optim="adamw_anyprecision"). Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_optimizer(name="adamw_torch", beta1=0.8)
>>> args.optim
'adamw_torch'

def set_push_to_hub(self,
model_id: str,
strategy: transformers.trainer_utils.HubStrategy | str = 'every_save',
token: str | None = None,
private_repo: bool | None = None,
always_push: bool = False,
revision: str | None = None)

Expand source code

def set_push_to_hub(
    self,
    model_id: str,
    strategy: str | HubStrategy = "every_save",
    token: str | None = None,
    private_repo: bool | None = None,
    always_push: bool = False,
    revision: str | None = None,
):
    """
    A method that regroups all arguments linked to synchronizing checkpoints with the Hub.

    <Tip>

    Calling this method will set `self.push_to_hub` to `True`, which means the `output_dir` will begin a git
    directory synced with the repo (determined by `model_id`) and the content will be pushed each time a save is
    triggered (depending on your `self.save_strategy`). Calling [`~Trainer.save_model`] will also trigger a push.

    </Tip>

    Args:
        model_id (`str`):
            The name of the repository to keep in sync with the local *output_dir*. It can be a simple model ID in
            which case the model will be pushed in your namespace. Otherwise it should be the whole repository
            name, for instance `"user_name/model"`, which allows you to push to an organization you are a member of
            with `"organization_name/model"`.
        strategy (`str` or [`~trainer_utils.HubStrategy`], *optional*, defaults to `"every_save"`):
            Defines the scope of what is pushed to the Hub and when. Possible values are:

            - `"end"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`]) and a
            draft of a model card when the [`~Trainer.save_model`] method is called.
            - `"every_save"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`])
              and
            a draft of a model card each time there is a model save. The pushes are asynchronous to not block
            training, and in case the save are very frequent, a new push is only attempted if the previous one is
            finished. A last push is made with the final model at the end of training.
            - `"checkpoint"`: like `"every_save"` but the latest checkpoint is also pushed in a subfolder named
            last-checkpoint, allowing you to resume training easily with
            `trainer.train(resume_from_checkpoint="last-checkpoint")`.
            - `"all_checkpoints"`: like `"checkpoint"` but all checkpoints are pushed like they appear in the
              output
            folder (so you will get one checkpoint folder per folder in your final repository)

        token (`str`, *optional*):
            The token to use to push the model to the Hub. Will default to the token in the cache folder obtained
            with `hf auth login`.
        private_repo (`bool`, *optional*, defaults to `False`):
            Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.
        always_push (`bool`, *optional*, defaults to `False`):
            Unless this is `True`, the `Trainer` will skip pushing a checkpoint when the previous push is not
            finished.
        revision (`str`, *optional*):
            The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash.

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_push_to_hub("me/awesome-model")
    >>> args.hub_model_id
    'me/awesome-model'
    ```
    """
    self.push_to_hub = True
    self.hub_model_id = model_id
    self.hub_strategy = HubStrategy(strategy)
    self.hub_token = token
    self.hub_private_repo = private_repo
    self.hub_always_push = always_push
    self.hub_revision = revision
    return self

A method that regroups all arguments linked to synchronizing checkpoints with the Hub.

Calling this method will set self.push_to_hub to True, which means the output_dir will begin a git directory synced with the repo (determined by model_id) and the content will be pushed each time a save is triggered (depending on your self.save_strategy). Calling [~Trainer.save_model] will also trigger a push.

Args -----= model_id (str): The name of the repository to keep in sync with the local output_dir. It can be a simple model ID in which case the model will be pushed in your namespace. Otherwise it should be the whole repository name, for instance "user_name/model", which allows you to push to an organization you are a member of with "organization_name/model". strategy (str or [~trainer_utils.HubStrategy], optional, defaults to "every_save"): Defines the scope of what is pushed to the Hub and when. Possible values are:

- `"end"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`]) and a
draft of a model card when the [`~Trainer.save_model`] method is called.
- `"every_save"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`])
  and
a draft of a model card each time there is a model save. The pushes are asynchronous to not block
training, and in case the save are very frequent, a new push is only attempted if the previous one is
finished. A last push is made with the final model at the end of training.
- `"checkpoint"`: like `"every_save"` but the latest checkpoint is also pushed in a subfolder named
last-checkpoint, allowing you to resume training easily with
`trainer.train(resume_from_checkpoint="last-checkpoint")`.
- `"all_checkpoints"`: like `"checkpoint"` but all checkpoints are pushed like they appear in the
  output
folder (so you will get one checkpoint folder per folder in your final repository)

token (str, optional): The token to use to push the model to the Hub. Will default to the token in the cache folder obtained with hf auth login. private_repo (bool, optional, defaults to False): Whether to make the repo private. If None (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. always_push (bool, optional, defaults to False): Unless this is True, the Trainer will skip pushing a checkpoint when the previous push is not finished. revision (str, optional): The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash. Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_push_to_hub("me/awesome-model")
>>> args.hub_model_id
'me/awesome-model'

def set_save(self,
strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps',
steps: int = 500,
total_limit: int | None = None,
on_each_node: bool = False)

Expand source code

def set_save(
    self,
    strategy: str | IntervalStrategy = "steps",
    steps: int = 500,
    total_limit: int | None = None,
    on_each_node: bool = False,
):
    """
    A method that regroups all arguments linked to checkpoint saving.

    Args:
        strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`):
            The checkpoint save strategy to adopt during training. Possible values are:

                - `"no"`: No save is done during training.
                - `"epoch"`: Save is done at the end of each epoch.
                - `"steps"`: Save is done every `save_steps`.

        steps (`int`, *optional*, defaults to 500):
            Number of updates steps before two checkpoint saves if `strategy="steps"`.
        total_limit (`int`, *optional*):
            If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in
            `output_dir`.
        on_each_node (`bool`, *optional*, defaults to `False`):
            When doing multi-node distributed training, whether to save models and checkpoints on each node, or
            only on the main one.

            This should not be activated when the different nodes use the same storage as the files will be saved
            with the same names for each node.

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_save(strategy="steps", steps=100)
    >>> args.save_steps
    100
    ```
    """
    self.save_strategy = SaveStrategy(strategy)
    if self.save_strategy == SaveStrategy.STEPS and steps == 0:
        raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.")
    self.save_steps = steps
    self.save_total_limit = total_limit
    self.save_on_each_node = on_each_node
    return self

A method that regroups all arguments linked to checkpoint saving.

Args -----= strategy (str or [~trainer_utils.IntervalStrategy], optional, defaults to "steps"): The checkpoint save strategy to adopt during training. Possible values are:

    - `"no"`: No save is done during training.
    - `"epoch"`: Save is done at the end of each epoch.
    - `"steps"`: Save is done every <code>save\_steps</code>.

steps (int, optional, defaults to 500): Number of updates steps before two checkpoint saves if strategy="steps". total_limit (int, optional): If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in output_dir. on_each_node (bool, optional, defaults to False): When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on the main one.

This should not be activated when the different nodes use the same storage as the files will be saved
with the same names for each node.

Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_save(strategy="steps", steps=100)
>>> args.save_steps
100

def set_testing(self, batch_size: int = 8, loss_only: bool = False)

Expand source code

def set_testing(
    self,
    batch_size: int = 8,
    loss_only: bool = False,
):
    """
    A method that regroups all basic arguments linked to testing on a held-out dataset.

    <Tip>

    Calling this method will automatically set `self.do_predict` to `True`.

    </Tip>

    Args:
        batch_size (`int` *optional*, defaults to 8):
            The batch size per device (GPU/TPU core/CPU...) used for testing.
        loss_only (`bool`, *optional*, defaults to `False`):
            Ignores all outputs except the loss.

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_testing(batch_size=32)
    >>> args.per_device_eval_batch_size
    32
    ```
    """
    self.do_predict = True
    self.per_device_eval_batch_size = batch_size
    self.prediction_loss_only = loss_only
    return self

A method that regroups all basic arguments linked to testing on a held-out dataset.

Calling this method will automatically set self.do_predict to True.

Args -----= batch_size (int optional, defaults to 8): The batch size per device (GPU/TPU core/CPU…) used for testing. loss_only (bool, optional, defaults to False): Ignores all outputs except the loss. Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_testing(batch_size=32)
>>> args.per_device_eval_batch_size
32

def set_training(self,
learning_rate: float = 5e-05,
batch_size: int = 8,
weight_decay: float = 0,
num_epochs: float = 3,
max_steps: int = -1,
gradient_accumulation_steps: int = 1,
seed: int = 42,
gradient_checkpointing: bool = False)

Expand source code

def set_training(
    self,
    learning_rate: float = 5e-5,
    batch_size: int = 8,
    weight_decay: float = 0,
    num_epochs: float = 3,
    max_steps: int = -1,
    gradient_accumulation_steps: int = 1,
    seed: int = 42,
    gradient_checkpointing: bool = False,
):
    """
    A method that regroups all basic arguments linked to the training.

    <Tip>

    Calling this method will automatically set `self.do_train` to `True`.

    </Tip>

    Args:
        learning_rate (`float`, *optional*, defaults to 5e-5):
            The initial learning rate for the optimizer.
        batch_size (`int` *optional*, defaults to 8):
            The batch size per device (GPU/TPU core/CPU...) used for training.
        weight_decay (`float`, *optional*, defaults to 0):
            The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights in the
            optimizer.
        num_train_epochs(`float`, *optional*, defaults to 3.0):
            Total number of training epochs to perform (if not an integer, will perform the decimal part percents
            of the last epoch before stopping training).
        max_steps (`int`, *optional*, defaults to -1):
            If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`.
            For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until
            `max_steps` is reached.
        gradient_accumulation_steps (`int`, *optional*, defaults to 1):
            Number of updates steps to accumulate the gradients for, before performing a backward/update pass.

            <Tip warning={true}>

            When using gradient accumulation, one step is counted as one step with backward pass. Therefore,
            logging, evaluation, save will be conducted every `gradient_accumulation_steps * xxx_step` training
            examples.

            </Tip>

        seed (`int`, *optional*, defaults to 42):
            Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use
            the [`~Trainer.model_init`] function to instantiate the model if it has some randomly initialized
            parameters.
        gradient_checkpointing (`bool`, *optional*, defaults to `False`):
            If True, use gradient checkpointing to save memory at the expense of slower backward pass.

    Example:

    ```py
    >>> from transformers import TrainingArguments

    >>> args = TrainingArguments("working_dir")
    >>> args = args.set_training(learning_rate=1e-4, batch_size=32)
    >>> args.learning_rate
    1e-4
    ```
    """
    self.do_train = True
    self.learning_rate = learning_rate
    self.per_device_train_batch_size = batch_size
    self.weight_decay = weight_decay
    self.num_train_epochs = num_epochs
    self.max_steps = max_steps
    self.gradient_accumulation_steps = gradient_accumulation_steps
    self.seed = seed
    self.gradient_checkpointing = gradient_checkpointing
    return self

A method that regroups all basic arguments linked to the training.

Calling this method will automatically set self.do_train to True.

Args -----= learning_rate (float, optional, defaults to 5e-5): The initial learning rate for the optimizer. batch_size (int optional, defaults to 8): The batch size per device (GPU/TPU core/CPU…) used for training. weight_decay (float, optional, defaults to 0): The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights in the optimizer. num_train_epochs(float, optional, defaults to 3.0): Total number of training epochs to perform (if not an integer, will perform the decimal part percents of the last epoch before stopping training). max_steps (int, optional, defaults to -1): If set to a positive number, the total number of training steps to perform. Overrides num_train_epochs. For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until max_steps is reached. gradient_accumulation_steps (int, optional, defaults to 1): Number of updates steps to accumulate the gradients for, before performing a backward/update pass.

<Tip warning={true}>

When using gradient accumulation, one step is counted as one step with backward pass. Therefore,
logging, evaluation, save will be conducted every `gradient_accumulation_steps * xxx_step` training
examples.

</Tip>

seed (int, optional, defaults to 42): Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use the [~Trainer.model_init] function to instantiate the model if it has some randomly initialized parameters. gradient_checkpointing (bool, optional, defaults to False): If True, use gradient checkpointing to save memory at the expense of slower backward pass. Example:

>>> from transformers import TrainingArguments

>>> args = TrainingArguments("working_dir")
>>> args = args.set_training(learning_rate=1e-4, batch_size=32)
>>> args.learning_rate
1e-4

def to_dict(self)

Expand source code

def to_dict(self):
    """
    Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates
    the token values by removing their value.
    """
    # Exclude non-init fields (they aren't user-facing config)
    d = {field.name: getattr(self, field.name) for field in fields(self) if field.init}

    for k, v in d.items():
        if isinstance(v, Enum):
            d[k] = v.value
        if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum):
            d[k] = [x.value for x in v]
        if k.endswith("_token"):
            d[k] = f"<{k.upper()}>"
        # Serialize AcceleratorConfig to dict
        if is_accelerate_available() and isinstance(v, AcceleratorConfig):
            d[k] = v.to_dict()
        # Serialize quantization_config if nested inside model_init_kwargs
        if k == "model_init_kwargs" and isinstance(v, dict) and "quantization_config" in v:
            quantization_config = v.get("quantization_config")
            if quantization_config and not isinstance(quantization_config, dict):
                d[k]["quantization_config"] = quantization_config.to_dict()
        if k == "parallelism_config" and v is not None:
            d[k] = v.to_json()

    self._dict_dtype_to_str(d)

    return d

Serializes this instance while replace Enum by their values (for JSON serialization support). It obfuscates the token values by removing their value.

def to_json_string(self)

Expand source code

def to_json_string(self):
    """
    Serializes this instance to a JSON string.
    """
    return json.dumps(self.to_dict(), indent=2)

Serializes this instance to a JSON string.

def to_sanitized_dict(self) ‑> dict[str, typing.Any]

Expand source code

def to_sanitized_dict(self) -> dict[str, Any]:
    """
    Sanitized serialization to use with TensorBoard's hparams
    """
    d = self.to_dict()
    d = {**d, "train_batch_size": self.train_batch_size, "eval_batch_size": self.eval_batch_size}

    valid_types = [bool, int, float, str]
    if is_torch_available():
        valid_types.append(torch.Tensor)

    return {k: v if type(v) in valid_types else str(v) for k, v in d.items()}

Sanitized serialization to use with TensorBoard's hparams

class TritonCosformerAttention (config: Any = None,
embed_dim: int | None = None,
num_heads: int | None = None,
num_kv_heads: int | None = None,
is_decoder: bool = False,
eps: float = 1e-06,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("triton_cosformer", supports_causal=True, supports_bidirectional=True)
class TritonCosformerAttention(BaseAttention):
    r"""
    Triton-Accelerated Cosformer Linear Attention with Full Autograd Differentiability,
    RoPE support, Pretrained Weight Loading, and Masked Attention.
    """

    triton_kernel = "linear_attn"

    def __init__(
        self,
        config: Any = None,
        embed_dim: int | None = None,
        num_heads: int | None = None,
        num_kv_heads: int | None = None,
        is_decoder: bool = False,
        eps: float = 1e-6,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ):
        super().__init__()
        self.config = config
        self.is_decoder = bool(
            is_decoder
            or getattr(config, "is_decoder", False)
            or getattr(config, "is_causal", False)
        )
        self.eps = eps
        self.triton_kwargs = {
            "phi": "elu+1",
            "eps": self.eps,
            "cos_weight": None,
            "sin_weight": None,
        }

        self.embed_dim = (
            embed_dim or kwargs.get("hidden_size") or getattr(config, "hidden_size", 768)
        )
        self.num_heads = (
            num_heads
            or kwargs.get("num_attention_heads")
            or getattr(config, "num_attention_heads", 12)
        )
        self.num_kv_heads = (
            num_kv_heads
            or kwargs.get("num_kv_heads")
            or getattr(config, "num_key_value_heads", self.num_heads)
        )
        self.head_dim = kwargs.get("head_dim") or (self.embed_dim // self.num_heads)

        # BUG 5 FIX: fixed max_length from config
        self.max_length = getattr(config, "max_position_embeddings", 2048) or 2048

        # Linear projections
        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.kv_dim = self.num_kv_heads * self.head_dim
        self.k_proj = nn.Linear(self.embed_dim, self.kv_dim, bias=True)
        self.v_proj = nn.Linear(self.embed_dim, self.kv_dim, bias=True)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)

        for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
            nn.init.xavier_uniform_(m.weight)
            if m.bias is not None:
                nn.init.zeros_(m.bias)

        # BUG 3 FIX: Support pretrained attention weight copying during initialization
        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    def copy_weights(self, src: Any) -> None:
        """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers."""
        if src is None:
            return

        getattr(src, "__class__", type(src)).__name__

        # 1. GPT-2: Fused c_attn (Conv1D) + c_proj
        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()
                # If Conv1D or shape is [in, out], transpose to standard PyTorch [out, in]
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or hasattr(c_attn_mod, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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

                self.q_proj.weight.data.copy_(c_w[:split_size, :])
                self.k_proj.weight.data.copy_(c_w[split_size : 2 * split_size, :])
                self.v_proj.weight.data.copy_(c_w[2 * split_size :, :])

                if c_b is not None:
                    if self.q_proj.bias is not None:
                        self.q_proj.bias.data.copy_(c_b[:split_size])
                    if self.k_proj.bias is not None:
                        self.k_proj.bias.data.copy_(c_b[split_size : 2 * split_size])
                    if self.v_proj.bias is not None:
                        self.v_proj.bias.data.copy_(c_b[2 * split_size :])

            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()
                    if (
                        c_proj_mod.__class__.__name__ == "Conv1D"
                        or base_o.__class__.__name__ == "Conv1D"
                        or hasattr(c_proj_mod, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        o_w = o_w.t().contiguous()
                    self.out_proj.weight.data.copy_(o_w)
                    bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                    if bias_o is not None and self.out_proj.bias is not None:
                        self.out_proj.bias.data.copy_(bias_o.data)

        # 2. RoBERTa / BERT nested self-attention
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")]
            for target_name, sub_name in mapping:
                target_mod = getattr(self, target_name)
                src_mod = getattr(s, sub_name, None)
                base_src = getattr(src_mod, "base_layer", src_mod)
                if base_src is not None and hasattr(base_src, "weight"):
                    target_mod.weight.data.copy_(base_src.weight.data)
                    if base_src.bias is not None and target_mod.bias is not None:
                        target_mod.bias.data.copy_(base_src.bias.data)

            if hasattr(src, "output") and hasattr(src.output, "dense"):
                out_mod = src.output.dense
                base_out = getattr(out_mod, "base_layer", out_mod)
                self.out_proj.weight.data.copy_(base_out.weight.data)
                if base_out.bias is not None and self.out_proj.bias is not None:
                    self.out_proj.bias.data.copy_(base_out.bias.data)

        # 3. Standard flat linear projections (DistilBERT, Qwen, LLaMA)
        else:
            proj_candidates = {
                "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"],
                "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"],
                "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"],
                "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"],
            }
            for target_name, candidates in proj_candidates.items():
                target_mod = getattr(self, target_name)
                for cand in candidates:
                    if hasattr(src, cand):
                        src_mod = getattr(src, cand)
                        base_mod = getattr(src_mod, "base_layer", src_mod)
                        if base_mod is not None and hasattr(base_mod, "weight"):
                            w = base_mod.weight.data
                            if (
                                w.ndim == 2
                                and w.shape != target_mod.weight.shape
                                and w.t().shape == target_mod.weight.shape
                            ):
                                w = w.t()
                            target_mod.weight.data.copy_(w)
                            if (
                                getattr(base_mod, "bias", None) is not None
                                and target_mod.bias is not None
                            ):
                                target_mod.bias.data.copy_(base_mod.bias.data)
                            break

    # Alias for API uniformity
    copy_weights_from_pretrained = copy_weights

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        position_ids: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, Any | None]:
        B, N, C = hidden_states.shape

        # RoPE from kwargs if not explicit
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")

        # 1. Project Q, K, V
        q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. BUG 4 FIX: Apply Rotary Position Embeddings (RoPE) if provided
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=position_ids)

        # 3. GQA expansion if num_kv_heads < num_heads
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.num_heads // self.num_kv_heads)
            v = repeat_kv(v, self.num_heads // self.num_kv_heads)

        # 4. BUG 2 FIX: Mask handling for padding tokens on K and V
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                # Shape (B, N) -> (B, 1, N, 1)
                mask = attention_mask[:, None, :, None].to(q.dtype)
                if (mask < 0).any():
                    mask = (mask > -1000.0).to(q.dtype)
                k = k * mask
                v = v * mask
            elif attention_mask.dim() == 4:
                if attention_mask.shape[-2] == 1:
                    mask = attention_mask.squeeze(1).squeeze(1)[:, None, :, None].to(q.dtype)
                else:
                    mask = attention_mask[:, 0, -1, :][:, None, :, None].to(q.dtype)
                if (mask < 0).any():
                    mask = (mask > -1000.0).to(q.dtype)
                k = k * mask
                v = v * mask

        # 5. Differentiable forward pass using cosformer's own autograd function.
        # NOTE: We call _TritonCosformerAutogradFunction.apply() directly instead
        # of routing through triton_dispatch → linear_attn.fwd, because:
        #   - linear_attn.fwd early-exits to PyTorch fallback when cos_weight is set
        #   - That path was broken (ctx.save_for_backward not called before return)
        #   - This module has its own complete Triton + PyTorch fallback built in.
        out = _TritonCosformerAutogradFunction.apply(
            q, k, v, self.is_decoder, self.eps, self.max_length
        )

        out = out.transpose(1, 2).contiguous().view(B, N, C)
        out = self.out_proj(out)
        return out, None

Triton-Accelerated Cosformer Linear Attention with Full Autograd Differentiability, RoPE support, Pretrained Weight Loading, and Masked Attention.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• BaseAttention
• torch.nn.modules.module.Module

Class variables

var triton_kernel

The type of the None singleton.

Methods

def copy_weights(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers."""
    if src is None:
        return

    getattr(src, "__class__", type(src)).__name__

    # 1. GPT-2: Fused c_attn (Conv1D) + c_proj
    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()
            # If Conv1D or shape is [in, out], transpose to standard PyTorch [out, in]
            if (
                c_attn_mod.__class__.__name__ == "Conv1D"
                or base_c.__class__.__name__ == "Conv1D"
                or hasattr(c_attn_mod, "nf")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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

            self.q_proj.weight.data.copy_(c_w[:split_size, :])
            self.k_proj.weight.data.copy_(c_w[split_size : 2 * split_size, :])
            self.v_proj.weight.data.copy_(c_w[2 * split_size :, :])

            if c_b is not None:
                if self.q_proj.bias is not None:
                    self.q_proj.bias.data.copy_(c_b[:split_size])
                if self.k_proj.bias is not None:
                    self.k_proj.bias.data.copy_(c_b[split_size : 2 * split_size])
                if self.v_proj.bias is not None:
                    self.v_proj.bias.data.copy_(c_b[2 * split_size :])

        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()
                if (
                    c_proj_mod.__class__.__name__ == "Conv1D"
                    or base_o.__class__.__name__ == "Conv1D"
                    or hasattr(c_proj_mod, "nf")
                    or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                ):
                    o_w = o_w.t().contiguous()
                self.out_proj.weight.data.copy_(o_w)
                bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                if bias_o is not None and self.out_proj.bias is not None:
                    self.out_proj.bias.data.copy_(bias_o.data)

    # 2. RoBERTa / BERT nested self-attention
    elif hasattr(src, "self") and hasattr(src.self, "query"):
        s = src.self
        mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")]
        for target_name, sub_name in mapping:
            target_mod = getattr(self, target_name)
            src_mod = getattr(s, sub_name, None)
            base_src = getattr(src_mod, "base_layer", src_mod)
            if base_src is not None and hasattr(base_src, "weight"):
                target_mod.weight.data.copy_(base_src.weight.data)
                if base_src.bias is not None and target_mod.bias is not None:
                    target_mod.bias.data.copy_(base_src.bias.data)

        if hasattr(src, "output") and hasattr(src.output, "dense"):
            out_mod = src.output.dense
            base_out = getattr(out_mod, "base_layer", out_mod)
            self.out_proj.weight.data.copy_(base_out.weight.data)
            if base_out.bias is not None and self.out_proj.bias is not None:
                self.out_proj.bias.data.copy_(base_out.bias.data)

    # 3. Standard flat linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"],
            "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"],
            "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"],
            "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"],
        }
        for target_name, candidates in proj_candidates.items():
            target_mod = getattr(self, target_name)
            for cand in candidates:
                if hasattr(src, cand):
                    src_mod = getattr(src, cand)
                    base_mod = getattr(src_mod, "base_layer", src_mod)
                    if base_mod is not None and hasattr(base_mod, "weight"):
                        w = base_mod.weight.data
                        if (
                            w.ndim == 2
                            and w.shape != target_mod.weight.shape
                            and w.t().shape == target_mod.weight.shape
                        ):
                            w = w.t()
                        target_mod.weight.data.copy_(w)
                        if (
                            getattr(base_mod, "bias", None) is not None
                            and target_mod.bias is not None
                        ):
                            target_mod.bias.data.copy_(base_mod.bias.data)
                        break

Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.

def copy_weights_from_pretrained(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers."""
    if src is None:
        return

    getattr(src, "__class__", type(src)).__name__

    # 1. GPT-2: Fused c_attn (Conv1D) + c_proj
    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()
            # If Conv1D or shape is [in, out], transpose to standard PyTorch [out, in]
            if (
                c_attn_mod.__class__.__name__ == "Conv1D"
                or base_c.__class__.__name__ == "Conv1D"
                or hasattr(c_attn_mod, "nf")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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

            self.q_proj.weight.data.copy_(c_w[:split_size, :])
            self.k_proj.weight.data.copy_(c_w[split_size : 2 * split_size, :])
            self.v_proj.weight.data.copy_(c_w[2 * split_size :, :])

            if c_b is not None:
                if self.q_proj.bias is not None:
                    self.q_proj.bias.data.copy_(c_b[:split_size])
                if self.k_proj.bias is not None:
                    self.k_proj.bias.data.copy_(c_b[split_size : 2 * split_size])
                if self.v_proj.bias is not None:
                    self.v_proj.bias.data.copy_(c_b[2 * split_size :])

        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()
                if (
                    c_proj_mod.__class__.__name__ == "Conv1D"
                    or base_o.__class__.__name__ == "Conv1D"
                    or hasattr(c_proj_mod, "nf")
                    or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                ):
                    o_w = o_w.t().contiguous()
                self.out_proj.weight.data.copy_(o_w)
                bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                if bias_o is not None and self.out_proj.bias is not None:
                    self.out_proj.bias.data.copy_(bias_o.data)

    # 2. RoBERTa / BERT nested self-attention
    elif hasattr(src, "self") and hasattr(src.self, "query"):
        s = src.self
        mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")]
        for target_name, sub_name in mapping:
            target_mod = getattr(self, target_name)
            src_mod = getattr(s, sub_name, None)
            base_src = getattr(src_mod, "base_layer", src_mod)
            if base_src is not None and hasattr(base_src, "weight"):
                target_mod.weight.data.copy_(base_src.weight.data)
                if base_src.bias is not None and target_mod.bias is not None:
                    target_mod.bias.data.copy_(base_src.bias.data)

        if hasattr(src, "output") and hasattr(src.output, "dense"):
            out_mod = src.output.dense
            base_out = getattr(out_mod, "base_layer", out_mod)
            self.out_proj.weight.data.copy_(base_out.weight.data)
            if base_out.bias is not None and self.out_proj.bias is not None:
                self.out_proj.bias.data.copy_(base_out.bias.data)

    # 3. Standard flat linear projections (DistilBERT, Qwen, LLaMA)
    else:
        proj_candidates = {
            "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"],
            "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"],
            "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"],
            "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"],
        }
        for target_name, candidates in proj_candidates.items():
            target_mod = getattr(self, target_name)
            for cand in candidates:
                if hasattr(src, cand):
                    src_mod = getattr(src, cand)
                    base_mod = getattr(src_mod, "base_layer", src_mod)
                    if base_mod is not None and hasattr(base_mod, "weight"):
                        w = base_mod.weight.data
                        if (
                            w.ndim == 2
                            and w.shape != target_mod.weight.shape
                            and w.t().shape == target_mod.weight.shape
                        ):
                            w = w.t()
                        target_mod.weight.data.copy_(w)
                        if (
                            getattr(base_mod, "bias", None) is not None
                            and target_mod.bias is not None
                        ):
                            target_mod.bias.data.copy_(base_mod.bias.data)
                        break

Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class TritonSVDLinearAttention (config: Any = None,
embed_dim: int | None = None,
num_heads: int | None = None,
num_kv_heads: int | None = None,
rank: int | None = 64,
is_decoder: bool = False,
activation_mode: str = 'elu',
eps: float = 1e-06,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "triton_svd_linear_attention", supports_causal=True, supports_bidirectional=True
)
@ATTENTION_REGISTRY.register(
    "triton_linear_attention", supports_causal=True, supports_bidirectional=True
)
class TritonSVDLinearAttention(BaseAttention):
    r"""
    Triton-Accelerated Linear Attention with Genuine SVD Low-Rank Projections,
    RoPE, GQA, and Masked Attention.
    """

    triton_kernel = "linear_attn"

    def __init__(
        self,
        config: Any = None,
        embed_dim: int | None = None,
        num_heads: int | None = None,
        num_kv_heads: int | None = None,
        rank: int | None = 64,
        is_decoder: bool = False,
        activation_mode: str = "elu",
        eps: float = 1e-6,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ):
        super().__init__()
        self.config = config
        self.is_decoder = bool(
            is_decoder
            or getattr(config, "is_decoder", False)
            or getattr(config, "is_causal", False)
        )
        self.activation_mode = activation_mode
        self.eps = eps
        self.triton_kwargs = {
            "phi": "elu+1" if self.activation_mode == "elu" else self.activation_mode,
            "eps": self.eps,
            "cos_weight": None,
            "sin_weight": None,
        }

        # Dimensionality extraction
        self.embed_dim = (
            embed_dim or kwargs.get("hidden_size") or getattr(config, "hidden_size", 768)
        )
        self.num_heads = (
            num_heads
            or kwargs.get("num_attention_heads")
            or getattr(config, "num_attention_heads", 12)
        )
        # BUG 4 FIX: Support GQA
        self.num_kv_heads = (
            num_kv_heads
            or kwargs.get("num_kv_heads")
            or getattr(config, "num_key_value_heads", self.num_heads)
        )
        self.head_dim = kwargs.get("head_dim") or (self.embed_dim // self.num_heads)
        self.rank = rank or 64
        self.kv_dim = self.num_kv_heads * self.head_dim

        # Genuine SVD Low-Rank Linear Projections
        if self.rank < self.embed_dim:
            self.q_proj = SVDLinearProjection(
                self.embed_dim, self.embed_dim, rank=self.rank, bias=True
            )
            self.k_proj = SVDLinearProjection(
                self.embed_dim, self.kv_dim, rank=min(self.rank, self.kv_dim), bias=True
            )
            self.v_proj = SVDLinearProjection(
                self.embed_dim, self.kv_dim, rank=min(self.rank, self.kv_dim), bias=True
            )
        else:
            self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
            self.k_proj = nn.Linear(self.embed_dim, self.kv_dim, bias=True)
            self.v_proj = nn.Linear(self.embed_dim, self.kv_dim, bias=True)

        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    def copy_weights(self, src: Any) -> None:
        """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers."""
        if src is None:
            return

        getattr(src, "__class__", type(src)).__name__

        # 1. GPT-2: Fused c_attn (Conv1D) + c_proj
        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()
                if (
                    c_attn_mod.__class__.__name__ == "Conv1D"
                    or base_c.__class__.__name__ == "Conv1D"
                    or hasattr(c_attn_mod, "nf")
                    or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
                ):
                    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

                q_w, q_b = c_w[:split_size, :], (c_b[:split_size] if c_b is not None else None)
                k_w, k_b = (
                    c_w[split_size : 2 * split_size, :],
                    (c_b[split_size : 2 * split_size] if c_b is not None else None),
                )
                v_w, v_b = (
                    c_w[2 * split_size :, :],
                    (c_b[2 * split_size :] if c_b is not None else None),
                )

                self._copy_single_proj(self.q_proj, q_w, q_b)
                self._copy_single_proj(self.k_proj, k_w, k_b)
                self._copy_single_proj(self.v_proj, v_w, v_b)

            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()
                    if (
                        c_proj_mod.__class__.__name__ == "Conv1D"
                        or base_o.__class__.__name__ == "Conv1D"
                        or hasattr(c_proj_mod, "nf")
                        or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                    ):
                        o_w = o_w.t().contiguous()
                    self.out_proj.weight.data.copy_(o_w)
                    bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                    if bias_o is not None and self.out_proj.bias is not None:
                        self.out_proj.bias.data.copy_(bias_o.data)

        # 2. RoBERTa / BERT nested self-attention
        elif hasattr(src, "self") and hasattr(src.self, "query"):
            s = src.self
            mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")]
            for target_name, sub_name in mapping:
                target_mod = getattr(self, target_name)
                src_mod = getattr(s, sub_name, None)
                base_src = getattr(src_mod, "base_layer", src_mod)
                if base_src is not None and hasattr(base_src, "weight"):
                    w = base_src.weight.data
                    b = getattr(base_src, "bias", None)
                    b_data = b.data if b is not None else None
                    self._copy_single_proj(target_mod, w, b_data)

            if hasattr(src, "output") and hasattr(src.output, "dense"):
                out_mod = src.output.dense
                base_out = getattr(out_mod, "base_layer", out_mod)
                self.out_proj.weight.data.copy_(base_out.weight.data)
                if base_out.bias is not None and self.out_proj.bias is not None:
                    self.out_proj.bias.data.copy_(base_out.bias.data)

        # 3. Standard flat linear projections
        else:
            proj_candidates = {
                "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"],
                "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"],
                "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"],
                "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"],
            }
            for target_name, candidates in proj_candidates.items():
                target_mod = getattr(self, target_name)
                for cand in candidates:
                    if hasattr(src, cand):
                        src_mod = getattr(src, cand)
                        base_mod = getattr(src_mod, "base_layer", src_mod)
                        if base_mod is not None and hasattr(base_mod, "weight"):
                            w = base_mod.weight.data
                            b = getattr(base_mod, "bias", None)
                            b_data = b.data if b is not None else None
                            if target_name == "out_proj":
                                if (
                                    w.ndim == 2
                                    and w.shape != target_mod.weight.shape
                                    and w.t().shape == target_mod.weight.shape
                                ):
                                    w = w.t()
                                target_mod.weight.data.copy_(w)
                                if b_data is not None and target_mod.bias is not None:
                                    target_mod.bias.data.copy_(b_data)
                            else:
                                self._copy_single_proj(target_mod, w, b_data)
                            break

    def _copy_single_proj(
        self, target_mod: nn.Module, w: torch.Tensor, b: torch.Tensor | None
    ) -> None:
        """Decomposes weights into low-rank SVD factors or copies directly."""
        if isinstance(target_mod, SVDLinearProjection):
            with torch.no_grad():
                W_f = w.float()
                u, s, vh = torch.linalg.svd(W_f, full_matrices=False)
                r = target_mod.rank
                s_sqrt = torch.diag(torch.sqrt(s[:r]))
                target_mod.U.data.copy_(torch.matmul(u[:, :r], s_sqrt).to(w.dtype))
                target_mod.V.data.copy_(torch.matmul(s_sqrt, vh[:r, :]).to(w.dtype))
                if b is not None and target_mod.bias is not None:
                    target_mod.bias.data.copy_(b)
        else:
            if (
                w.ndim == 2
                and w.shape != target_mod.weight.shape
                and w.t().shape == target_mod.weight.shape
            ):
                w = w.t()
            target_mod.weight.data.copy_(w)
            if b is not None and target_mod.bias is not None:
                target_mod.bias.data.copy_(b)

    # Alias for API uniformity
    copy_weights_from_pretrained = copy_weights

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        position_ids: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, Any | None]:
        B, N, C = hidden_states.shape

        if torch.is_autocast_enabled() and hidden_states.is_cuda:
            ac_dtype = (
                torch.get_autocast_dtype("cuda")
                if hasattr(torch, "get_autocast_dtype")
                else torch.float16
            )
            hidden_states = hidden_states.to(ac_dtype)

        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")

        # 1. Project Q, K, V via SVD low-rank factorized projections
        q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, N, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # 2. BUG 3 FIX: Apply Rotary Position Embeddings (RoPE) if provided
        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=position_ids)

        # 3. BUG 4 FIX: GQA expansion if num_kv_heads < num_heads
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.num_heads // self.num_kv_heads)
            v = repeat_kv(v, self.num_heads // self.num_kv_heads)

        # 4. BUG 2 FIX: Mask handling for padding tokens on K and V
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                mask = attention_mask[:, None, :, None].to(q.dtype)  # [B, 1, N, 1]
                if (mask < 0).any():
                    mask = (mask > -1000.0).to(q.dtype)
                k = k * mask
                v = v * mask
            elif attention_mask.dim() == 4:
                if attention_mask.shape[-2] == 1:
                    mask = attention_mask.squeeze(1).squeeze(1)[:, None, :, None].to(q.dtype)
                else:
                    mask = attention_mask[:, 0, -1, :][:, None, :, None].to(q.dtype)
                if (mask < 0).any():
                    mask = (mask > -1000.0).to(q.dtype)
                k = k * mask
                v = v * mask

        # 5. Dispatch Fully Differentiable Triton Linear Attention
        out = triton_dispatch(
            self.triton_kernel,
            q,
            k,
            v,
            is_causal=self.is_decoder,
            **self.triton_kwargs,
        )

        # 6. Reshape and final output projection
        out = out.transpose(1, 2).contiguous().view(B, N, C)
        out = self.out_proj(out)

        # Numerical safety guard against unbounded linear attention residual accumulation in FP16
        if out.dtype == torch.float16:
            out = torch.clamp(out, min=-65000.0, max=65000.0)

        return out, None

Triton-Accelerated Linear Attention with Genuine SVD Low-Rank Projections, RoPE, GQA, and Masked Attention.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• BaseAttention
• torch.nn.modules.module.Module

Class variables

var triton_kernel

The type of the None singleton.

Methods

def copy_weights(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers."""
    if src is None:
        return

    getattr(src, "__class__", type(src)).__name__

    # 1. GPT-2: Fused c_attn (Conv1D) + c_proj
    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()
            if (
                c_attn_mod.__class__.__name__ == "Conv1D"
                or base_c.__class__.__name__ == "Conv1D"
                or hasattr(c_attn_mod, "nf")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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

            q_w, q_b = c_w[:split_size, :], (c_b[:split_size] if c_b is not None else None)
            k_w, k_b = (
                c_w[split_size : 2 * split_size, :],
                (c_b[split_size : 2 * split_size] if c_b is not None else None),
            )
            v_w, v_b = (
                c_w[2 * split_size :, :],
                (c_b[2 * split_size :] if c_b is not None else None),
            )

            self._copy_single_proj(self.q_proj, q_w, q_b)
            self._copy_single_proj(self.k_proj, k_w, k_b)
            self._copy_single_proj(self.v_proj, v_w, v_b)

        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()
                if (
                    c_proj_mod.__class__.__name__ == "Conv1D"
                    or base_o.__class__.__name__ == "Conv1D"
                    or hasattr(c_proj_mod, "nf")
                    or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                ):
                    o_w = o_w.t().contiguous()
                self.out_proj.weight.data.copy_(o_w)
                bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                if bias_o is not None and self.out_proj.bias is not None:
                    self.out_proj.bias.data.copy_(bias_o.data)

    # 2. RoBERTa / BERT nested self-attention
    elif hasattr(src, "self") and hasattr(src.self, "query"):
        s = src.self
        mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")]
        for target_name, sub_name in mapping:
            target_mod = getattr(self, target_name)
            src_mod = getattr(s, sub_name, None)
            base_src = getattr(src_mod, "base_layer", src_mod)
            if base_src is not None and hasattr(base_src, "weight"):
                w = base_src.weight.data
                b = getattr(base_src, "bias", None)
                b_data = b.data if b is not None else None
                self._copy_single_proj(target_mod, w, b_data)

        if hasattr(src, "output") and hasattr(src.output, "dense"):
            out_mod = src.output.dense
            base_out = getattr(out_mod, "base_layer", out_mod)
            self.out_proj.weight.data.copy_(base_out.weight.data)
            if base_out.bias is not None and self.out_proj.bias is not None:
                self.out_proj.bias.data.copy_(base_out.bias.data)

    # 3. Standard flat linear projections
    else:
        proj_candidates = {
            "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"],
            "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"],
            "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"],
            "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"],
        }
        for target_name, candidates in proj_candidates.items():
            target_mod = getattr(self, target_name)
            for cand in candidates:
                if hasattr(src, cand):
                    src_mod = getattr(src, cand)
                    base_mod = getattr(src_mod, "base_layer", src_mod)
                    if base_mod is not None and hasattr(base_mod, "weight"):
                        w = base_mod.weight.data
                        b = getattr(base_mod, "bias", None)
                        b_data = b.data if b is not None else None
                        if target_name == "out_proj":
                            if (
                                w.ndim == 2
                                and w.shape != target_mod.weight.shape
                                and w.t().shape == target_mod.weight.shape
                            ):
                                w = w.t()
                            target_mod.weight.data.copy_(w)
                            if b_data is not None and target_mod.bias is not None:
                                target_mod.bias.data.copy_(b_data)
                        else:
                            self._copy_single_proj(target_mod, w, b_data)
                        break

Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.

def copy_weights_from_pretrained(self,
src: Any) ‑> None

Expand source code

def copy_weights(self, src: Any) -> None:
    """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers."""
    if src is None:
        return

    getattr(src, "__class__", type(src)).__name__

    # 1. GPT-2: Fused c_attn (Conv1D) + c_proj
    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()
            if (
                c_attn_mod.__class__.__name__ == "Conv1D"
                or base_c.__class__.__name__ == "Conv1D"
                or hasattr(c_attn_mod, "nf")
                or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1])
            ):
                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

            q_w, q_b = c_w[:split_size, :], (c_b[:split_size] if c_b is not None else None)
            k_w, k_b = (
                c_w[split_size : 2 * split_size, :],
                (c_b[split_size : 2 * split_size] if c_b is not None else None),
            )
            v_w, v_b = (
                c_w[2 * split_size :, :],
                (c_b[2 * split_size :] if c_b is not None else None),
            )

            self._copy_single_proj(self.q_proj, q_w, q_b)
            self._copy_single_proj(self.k_proj, k_w, k_b)
            self._copy_single_proj(self.v_proj, v_w, v_b)

        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()
                if (
                    c_proj_mod.__class__.__name__ == "Conv1D"
                    or base_o.__class__.__name__ == "Conv1D"
                    or hasattr(c_proj_mod, "nf")
                    or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1])
                ):
                    o_w = o_w.t().contiguous()
                self.out_proj.weight.data.copy_(o_w)
                bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None))
                if bias_o is not None and self.out_proj.bias is not None:
                    self.out_proj.bias.data.copy_(bias_o.data)

    # 2. RoBERTa / BERT nested self-attention
    elif hasattr(src, "self") and hasattr(src.self, "query"):
        s = src.self
        mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")]
        for target_name, sub_name in mapping:
            target_mod = getattr(self, target_name)
            src_mod = getattr(s, sub_name, None)
            base_src = getattr(src_mod, "base_layer", src_mod)
            if base_src is not None and hasattr(base_src, "weight"):
                w = base_src.weight.data
                b = getattr(base_src, "bias", None)
                b_data = b.data if b is not None else None
                self._copy_single_proj(target_mod, w, b_data)

        if hasattr(src, "output") and hasattr(src.output, "dense"):
            out_mod = src.output.dense
            base_out = getattr(out_mod, "base_layer", out_mod)
            self.out_proj.weight.data.copy_(base_out.weight.data)
            if base_out.bias is not None and self.out_proj.bias is not None:
                self.out_proj.bias.data.copy_(base_out.bias.data)

    # 3. Standard flat linear projections
    else:
        proj_candidates = {
            "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"],
            "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"],
            "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"],
            "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"],
        }
        for target_name, candidates in proj_candidates.items():
            target_mod = getattr(self, target_name)
            for cand in candidates:
                if hasattr(src, cand):
                    src_mod = getattr(src, cand)
                    base_mod = getattr(src_mod, "base_layer", src_mod)
                    if base_mod is not None and hasattr(base_mod, "weight"):
                        w = base_mod.weight.data
                        b = getattr(base_mod, "bias", None)
                        b_data = b.data if b is not None else None
                        if target_name == "out_proj":
                            if (
                                w.ndim == 2
                                and w.shape != target_mod.weight.shape
                                and w.t().shape == target_mod.weight.shape
                            ):
                                w = w.t()
                            target_mod.weight.data.copy_(w)
                            if b_data is not None and target_mod.bias is not None:
                                target_mod.bias.data.copy_(b_data)
                        else:
                            self._copy_single_proj(target_mod, w, b_data)
                        break

Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class UnifiedMoE (dim: int,
num_experts: int = 2,
top_k: int = 1,
output_dim: int | None = None,
noise_std: float = 0.1,
expert_hidden_mult: float = 0.25,
gate_init: float = -3.0)

Expand source code

class UnifiedMoE(nn.Module):
    """
    Compact MoE with top-k gating.

    BAM-FIX-02 / FIX-12: expert hidden dim is `dim // (num_experts * 4)`,
    output dim is configurable (default `head_dim`, small).
    """

    is_custom_attention: bool = False

    def __init__(
        self,
        dim: int,
        num_experts: int = 2,
        top_k: int = 1,
        output_dim: int | None = None,
        noise_std: float = 0.1,
        expert_hidden_mult: float = 0.25,
        gate_init: float = -3.0,  # BAM-FIX-11
    ) -> None:
        super().__init__()
        self.dim = int(dim)
        self.num_experts = max(1, int(num_experts))
        self.top_k = min(max(1, int(top_k)), self.num_experts)
        # BAM-FIX-02: default to dim // 4 for compact output
        self.out_dim = int(output_dim) if output_dim is not None else max(16, self.dim // 4)
        self.noise_std = float(noise_std)

        self.router_ln = nn.LayerNorm(self.dim)
        self.router = nn.Linear(self.dim, self.num_experts, bias=False)
        nn.init.uniform_(self.router.weight, -0.01, 0.01)

        hidden_dim = max(16, int(self.out_dim * expert_hidden_mult))
        self.experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(self.dim, hidden_dim),
                    nn.GELU(),
                    nn.Linear(hidden_dim, self.out_dim),
                )
                for _ in range(self.num_experts)
            ]
        )

        for exp in self.experts:
            nn.init.normal_(exp[2].weight, std=0.02)
            if exp[2].bias is not None:
                nn.init.zeros_(exp[2].bias)

        # BAM-FIX-11: configurable gate init
        self.gate_bias = nn.Parameter(torch.tensor([float(gate_init)]))

        # BAM-FIX-04: aux loss is an attribute, not a tensor
        self._aux_loss: torch.Tensor | None = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape
        x_norm = self.router_ln(x)
        logits = self.router(x_norm)

        if self.training and self.noise_std > 0:
            logits = logits + torch.randn_like(logits) * self.noise_std

        all_probs = F.softmax(logits, dim=-1)
        topk_scores, topk_idx = torch.topk(all_probs, self.top_k, dim=-1)
        scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True).clamp(min=1e-6)
        scores = torch.nan_to_num(scores, nan=0.0)

        # Load balancing aux loss
        if self.training and self.num_experts > 1:
            router_probs = all_probs.view(-1, self.num_experts)
            expert_mask = F.one_hot(topk_idx, num_classes=self.num_experts).float()
            density = expert_mask.sum(dim=2).view(-1, self.num_experts).mean(dim=0)
            prob_mean = router_probs.mean(dim=0)
            self._aux_loss = self.num_experts * torch.sum(density * prob_mean)
        else:
            self._aux_loss = None

        experts_out = torch.stack([exp(x) for exp in self.experts], dim=2)
        topk_out = torch.gather(
            experts_out,
            2,
            topk_idx.unsqueeze(-1).expand(B, L, self.top_k, self.out_dim),
        )
        out = torch.sum(topk_out * scores.unsqueeze(-1), dim=2)
        out = torch.clamp(torch.nan_to_num(out, nan=0.0), -65000.0, 65000.0)
        gate = torch.sigmoid(self.gate_bias).to(x.dtype)
        return (out * gate).to(x.dtype)

    def get_aux_loss(self, device: torch.device) -> torch.Tensor:
        """BAM-FIX-04: return aux loss on the correct device."""
        if self._aux_loss is None:
            return torch.zeros((), device=device, dtype=torch.float32)
        return self._aux_loss.to(device)

Compact MoE with top-k gating.

BAM-FIX-02 / FIX-12: expert hidden dim is dim // (num_experts * 4), output dim is configurable (default head_dim, small).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    B, L, D = x.shape
    x_norm = self.router_ln(x)
    logits = self.router(x_norm)

    if self.training and self.noise_std > 0:
        logits = logits + torch.randn_like(logits) * self.noise_std

    all_probs = F.softmax(logits, dim=-1)
    topk_scores, topk_idx = torch.topk(all_probs, self.top_k, dim=-1)
    scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True).clamp(min=1e-6)
    scores = torch.nan_to_num(scores, nan=0.0)

    # Load balancing aux loss
    if self.training and self.num_experts > 1:
        router_probs = all_probs.view(-1, self.num_experts)
        expert_mask = F.one_hot(topk_idx, num_classes=self.num_experts).float()
        density = expert_mask.sum(dim=2).view(-1, self.num_experts).mean(dim=0)
        prob_mean = router_probs.mean(dim=0)
        self._aux_loss = self.num_experts * torch.sum(density * prob_mean)
    else:
        self._aux_loss = None

    experts_out = torch.stack([exp(x) for exp in self.experts], dim=2)
    topk_out = torch.gather(
        experts_out,
        2,
        topk_idx.unsqueeze(-1).expand(B, L, self.top_k, self.out_dim),
    )
    out = torch.sum(topk_out * scores.unsqueeze(-1), dim=2)
    out = torch.clamp(torch.nan_to_num(out, nan=0.0), -65000.0, 65000.0)
    gate = torch.sigmoid(self.gate_bias).to(x.dtype)
    return (out * gate).to(x.dtype)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def get_aux_loss(self, device: torch.device) ‑> torch.Tensor

Expand source code

def get_aux_loss(self, device: torch.device) -> torch.Tensor:
    """BAM-FIX-04: return aux loss on the correct device."""
    if self._aux_loss is None:
        return torch.zeros((), device=device, dtype=torch.float32)
    return self._aux_loss.to(device)

BAM-FIX-04: return aux loss on the correct device.

class UnifiedRoPE (config: Any | None = None,
rank: int | None = None,
dim: int | None = None,
num_heads: int | None = None,
max_seq_len: int = 4096,
base: int = 10000,
dropout: float = 0.1,
pretrained_attn: torch.nn.modules.module.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("compression_unified_rope", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("unified_rope", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("unified_rope_attention", supports_causal=True, supports_gqa=True)
class UnifiedRoPE(BaseAttention):
    r"""Universal Rotary Position Embedding Attention module."""

    supports_gqa: bool = True
    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        dim: int | None = None,
        num_heads: int | None = None,
        max_seq_len: int = 4096,
        base: int = 10000,
        dropout: float = 0.1,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config=config, rank=rank)

        src = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")

        if dim is None:
            if config is not None and not isinstance(config, int):
                dim = _cfg_get(
                    config,
                    "hidden_size",
                    _cfg_get(config, "dim", _cfg_get(config, "d_model", None)),
                )
            if dim is None and src is not None:
                dim = getattr(src, "hidden_size", getattr(src, "embed_dim", None))
                if dim is None and hasattr(src, "q_proj") and hasattr(src.q_proj, "in_features"):
                    dim = src.q_proj.in_features
                elif dim is None and hasattr(src, "q_lin") and hasattr(src.q_lin, "in_features"):
                    dim = src.q_lin.in_features
                elif dim is None and hasattr(src, "query") and hasattr(src.query, "in_features"):
                    dim = src.query.in_features
                elif dim is None and hasattr(src, "c_attn") and hasattr(src.c_attn, "weight"):
                    dim = (
                        src.c_attn.weight.shape[0]
                        if src.c_attn.weight.shape[0] != 3 * src.c_attn.weight.shape[1]
                        else src.c_attn.weight.shape[1]
                    )
            if dim is None:
                dim = kwargs.get("hidden_size", 768)

        if num_heads is None:
            if config is not None and not isinstance(config, int):
                num_heads = _cfg_get(
                    config,
                    "num_attention_heads",
                    _cfg_get(config, "n_heads", _cfg_get(config, "num_heads", None)),
                )
            if num_heads is None and src is not None:
                num_heads = getattr(
                    src,
                    "num_heads",
                    getattr(src, "num_attention_heads", getattr(src, "n_head", None)),
                )
            if num_heads is None:
                num_heads = 12

        self.hidden_size: int = int(dim)
        self.num_heads: int = max(1, int(num_heads))
        self.head_dim: int = self.hidden_size // self.num_heads
        self.max_seq_len: int = max_seq_len
        self.base: int = base
        self.scale: float = 1.0 / math.sqrt(self.head_dim)
        self.lora_rank: int = kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank))
        self.lora_alpha: float = kwargs.get(
            "lora_alpha", _cfg_get(config, "lora_alpha", lora_alpha)
        )

        cfg_kv = (
            _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
            if config is not None
            else None
        )
        if cfg_kv is None and src is not None:
            cfg_kv = getattr(src, "num_key_value_heads", getattr(src, "num_kv_heads", None))
        self.num_kv_heads: int = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_key_value_groups: int = max(1, self.num_heads // max(1, self.num_kv_heads))

        self.q_head_size: int = self.num_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim

        # Decoder / Causal flag
        if is_decoder is not None:
            self.is_decoder: bool = is_decoder
        elif config is not None:
            self.is_decoder = bool(_cfg_get(config, "is_decoder", False))
        else:
            self.is_decoder = False

        self.is_nested_encoder: bool = False

        self.q_proj = LoRALinear(
            self.hidden_size,
            self.q_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.q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        dropout_p = (
            getattr(
                config,
                "dropout",
                getattr(config, "attention_probs_dropout_prob", dropout),
            )
            if config
            else dropout
        )
        self.dropout = nn.Dropout(dropout_p)

        rotary_dim = 2 * (self.head_dim // 2)
        self.rotary_dim = rotary_dim
        if rotary_dim > 0:
            inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim))
        else:
            inv_freq = torch.empty(0)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._build_cache(max_seq_len)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    def _build_cache(self, seq_len: int) -> None:
        """Precomputes cosine and sine frequency tables up to seq_len."""
        if self.rotary_dim == 0 or self.inv_freq.numel() == 0:
            self.register_buffer(
                "cos_cached",
                torch.ones(seq_len, 0, device=self.inv_freq.device),
                persistent=False,
            )
            self.register_buffer(
                "sin_cached",
                torch.zeros(seq_len, 0, device=self.inv_freq.device),
                persistent=False,
            )
            return
        t = torch.arange(seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def _apply_rope(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        """Applies RoPE rotation: [B, H, L, D_head]."""
        if self.rotary_dim == 0 or self.cos_cached.shape[-1] == 0:
            return x

        if seq_len > self.max_seq_len:
            self._build_cache(seq_len)
            self.max_seq_len = seq_len

        rotary_dim = min(self.rotary_dim, self.cos_cached.shape[-1], x.shape[-1])
        rotary_dim = 2 * (rotary_dim // 2)
        if rotary_dim == 0:
            return x

        cos = self.cos_cached[:seq_len, :rotary_dim].unsqueeze(0).unsqueeze(0).to(x.device, x.dtype)
        sin = self.sin_cached[:seq_len, :rotary_dim].unsqueeze(0).unsqueeze(0).to(x.device, x.dtype)

        if rotary_dim < x.shape[-1]:
            x_rot = x[..., :rotary_dim]
            x_pass = x[..., rotary_dim:]
            rotated = (x_rot * cos) + (_rotate_half(x_rot) * sin)
            return torch.cat((rotated, x_pass), dim=-1)
        else:
            return (x * cos) + (_rotate_half(x) * sin)

    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        with torch.no_grad():
            w_q, b_q = None, None
            w_k, b_k = None, None
            w_v, b_v = None, None
            w_o, b_o = None, None

            # Check nested encoder container
            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight
                    b_o = old_attn.output.dense.bias
            else:
                src = old_attn

            # Pattern 1: GPT-2 Fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                if c_attn_w.shape[0] == self.hidden_size and c_attn_w.shape[1] != self.hidden_size:
                    c_attn_w = c_attn_w.t().contiguous()

                q_dim = self.q_head_size
                k_dim = self.kv_head_size
                v_dim = self.kv_head_size

                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

                out_w = src.c_proj.weight.detach()
                if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # Pattern 2: RoBERTa / BERT separate query, key, value
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None

                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None

                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None
                elif w_o is None:
                    for attr in (
                        "out_proj",
                        "out_lin",
                        "dense",
                        "c_proj",
                        "wo",
                        "proj",
                        "output",
                    ):
                        if hasattr(src, attr):
                            proj_layer = getattr(src, attr)
                            if hasattr(proj_layer, "dense"):
                                proj_layer = proj_layer.dense
                            if hasattr(proj_layer, "weight"):
                                w_o = proj_layer.weight.detach()
                                b_o = (
                                    proj_layer.bias.detach()
                                    if proj_layer.bias is not None
                                    else None
                                )
                                break

            # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None

                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None
                elif w_o is None:
                    for attr in (
                        "out_proj",
                        "o_proj",
                        "dense",
                        "c_proj",
                        "wo",
                        "proj",
                        "output",
                    ):
                        if hasattr(src, attr):
                            proj_layer = getattr(src, attr)
                            if hasattr(proj_layer, "dense"):
                                proj_layer = proj_layer.dense
                            if hasattr(proj_layer, "weight"):
                                w_o = proj_layer.weight.detach()
                                b_o = (
                                    proj_layer.bias.detach()
                                    if proj_layer.bias is not None
                                    else None
                                )
                                break

            # Pattern 5: UniversalAdapter fallback
            if w_q is None or w_k is None or w_v is None or w_o is None:
                try:
                    adapter = UniversalAdapter()
                    if w_q is None:
                        w_q, b_q = adapter.get_weight_bias(src, "q")
                    if w_k is None:
                        w_k, b_k = adapter.get_weight_bias(src, "k")
                    if w_v is None:
                        w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            # Strict validation
            missing = []
            if w_q is None:
                missing.append("Q projection")
            if w_k is None:
                missing.append("K projection")
            if w_v is None:
                missing.append("V projection")
            if w_o is None and not self.is_nested_encoder:
                missing.append("Out projection")

            if missing:
                raise RuntimeError(
                    f"UnifiedRoPE: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
                )

            def _copy_layer(target_layer, w, b, src_sub=None):
                if w is not None:
                    if w.shape == target_layer.weight.shape:
                        target_layer.weight.copy_(w.to(target_layer.weight.dtype))
                    elif w.t().contiguous().shape == target_layer.weight.shape:
                        target_layer.weight.copy_(w.t().contiguous().to(target_layer.weight.dtype))
                    else:
                        min_0 = min(target_layer.weight.shape[0], w.shape[0])
                        min_1 = min(target_layer.weight.shape[1], w.shape[1])
                        target_layer.weight.zero_()
                        target_layer.weight[:min_0, :min_1].copy_(
                            w[:min_0, :min_1].to(target_layer.weight.dtype)
                        )
                if b is not None and target_layer.bias is not None:
                    if b.shape == target_layer.bias.shape:
                        target_layer.bias.copy_(b.to(target_layer.bias.dtype))
                    else:
                        min_b = min(target_layer.bias.shape[0], b.shape[0])
                        target_layer.bias.zero_()
                        target_layer.bias[:min_b].copy_(b[:min_b].to(target_layer.bias.dtype))
                elif target_layer.bias is not None:
                    target_layer.bias.zero_()

                if (
                    src_sub is not None
                    and hasattr(src_sub, "lora_A")
                    and hasattr(src_sub, "lora_B")
                ):
                    if (
                        src_sub.lora_A is not None
                        and src_sub.lora_B is not None
                        and target_layer.lora_A is not None
                        and target_layer.lora_B is not None
                    ):
                        src_lA = getattr(src_sub.lora_A, "weight", src_sub.lora_A)
                        src_lB = getattr(src_sub.lora_B, "weight", src_sub.lora_B)
                        if hasattr(src_lA, "data") and hasattr(src_lB, "data"):
                            r_min = min(target_layer.lora_A.shape[0], src_lA.shape[0])
                            in_min = min(target_layer.lora_A.shape[1], src_lA.shape[1])
                            target_layer.lora_A.data[:r_min, :in_min].copy_(
                                src_lA.data[:r_min, :in_min]
                            )
                            target_layer.lora_A.requires_grad = True

                            out_min = min(target_layer.lora_B.shape[0], src_lB.shape[0])
                            r_min_b = min(target_layer.lora_B.shape[1], src_lB.shape[1])
                            target_layer.lora_B.data[:out_min, :r_min_b].copy_(
                                src_lB.data[:out_min, :r_min_b]
                            )
                            target_layer.lora_B.requires_grad = True
                        if hasattr(src_sub, "lora_scaling"):
                            target_layer.scaling = src_sub.lora_scaling

            sub_q = getattr(src, "q_proj", getattr(src, "q_lin", getattr(src, "query", None)))
            sub_k = getattr(src, "k_proj", getattr(src, "k_lin", getattr(src, "key", None)))
            sub_v = getattr(src, "v_proj", getattr(src, "v_lin", getattr(src, "value", None)))
            sub_o = getattr(
                src,
                "o_proj",
                getattr(
                    src,
                    "out_lin",
                    getattr(src, "dense", getattr(src, "out_proj", None)),
                ),
            )

            _copy_layer(self.q_proj, w_q, b_q, sub_q)
            _copy_layer(self.k_proj, w_k, b_k, sub_k)
            _copy_layer(self.v_proj, w_v, b_v, sub_v)
            if w_o is not None:
                _copy_layer(self.out_proj, w_o, b_o, sub_o)

                f_norm_o = float(torch.linalg.norm(self.out_proj.weight))
                if f_norm_o < 1e-4:
                    logger.warning(f"UnifiedRoPE: Out projection norm is near zero: {f_norm_o}")

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        """Forward pass computing RoPE attention, returning (output, attention_weights)."""
        B, L, _ = hidden_states.shape

        q = (
            self.q_proj(hidden_states)
            .view(B, L, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.k_proj(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.v_proj(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # Apply RoPE
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
            cos, sin = pos_emb
            q, k = apply_rotary_pos_emb(q, k, cos, sin)
        else:
            q = self._apply_rope(q, L)
            k = self._apply_rope(k, L)

        # GQA expansion
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.num_key_value_groups)
            v = repeat_kv(v, self.num_key_value_groups)

        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(q.float(), k.float().transpose(-1, -2).contiguous()) * scale

        # Causal mask for decoders
        if self.is_decoder and L > 1:
            causal_mask = torch.triu(
                torch.full((L, L), float("-inf"), device=scores.device, dtype=scores.dtype),
                diagonal=1,
            )
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

        mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask")
        if mask is not None:
            scores = apply_attention_mask(scores, mask)

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
        attn_weights = self.dropout(attn_weights.to(v.dtype))

        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:
                head_mask = head_mask
            attn_weights = attn_weights * head_mask.to(attn_weights.dtype)

        context = torch.matmul(attn_weights, v)
        context = context.transpose(1, 2).contiguous().view(B, L, self.num_heads * self.head_dim)

        if self.is_nested_encoder:
            if context.shape[-1] != self.hidden_size:
                output = self.out_proj(context)
            else:
                output = context
        else:
            output = self.out_proj(context)

        if output.shape[-1] != hidden_states.shape[-1]:
            output = (
                F.pad(output, (0, hidden_states.shape[-1] - output.shape[-1]))
                if output.shape[-1] < hidden_states.shape[-1]
                else output[..., : hidden_states.shape[-1]]
            )

        return (
            output.to(hidden_states.dtype),
            (attn_weights if output_attentions else None),
        )

Universal Rotary Position Embedding Attention module.

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 copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    with torch.no_grad():
        w_q, b_q = None, None
        w_k, b_k = None, None
        w_v, b_v = None, None
        w_o, b_o = None, None

        # Check nested encoder container
        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight
                b_o = old_attn.output.dense.bias
        else:
            src = old_attn

        # Pattern 1: GPT-2 Fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
            if c_attn_w.shape[0] == self.hidden_size and c_attn_w.shape[1] != self.hidden_size:
                c_attn_w = c_attn_w.t().contiguous()

            q_dim = self.q_head_size
            k_dim = self.kv_head_size
            v_dim = self.kv_head_size

            w_q = c_attn_w[:q_dim, :]
            w_k = c_attn_w[q_dim : q_dim + k_dim, :]
            w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]

            if c_attn_b is not None:
                b_q = c_attn_b[:q_dim]
                b_k = c_attn_b[q_dim : q_dim + k_dim]
                b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]

            out_w = src.c_proj.weight.detach()
            if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                out_w = out_w.t().contiguous()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # Pattern 2: RoBERTa / BERT separate query, key, value
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None

            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj)
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None

            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None
            elif w_o is None:
                for attr in (
                    "out_proj",
                    "out_lin",
                    "dense",
                    "c_proj",
                    "wo",
                    "proj",
                    "output",
                ):
                    if hasattr(src, attr):
                        proj_layer = getattr(src, attr)
                        if hasattr(proj_layer, "dense"):
                            proj_layer = proj_layer.dense
                        if hasattr(proj_layer, "weight"):
                            w_o = proj_layer.weight.detach()
                            b_o = (
                                proj_layer.bias.detach()
                                if proj_layer.bias is not None
                                else None
                            )
                            break

        # Pattern 4: DistilBERT (q_lin, k_lin, v_lin, out_lin)
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None

            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None
            elif w_o is None:
                for attr in (
                    "out_proj",
                    "o_proj",
                    "dense",
                    "c_proj",
                    "wo",
                    "proj",
                    "output",
                ):
                    if hasattr(src, attr):
                        proj_layer = getattr(src, attr)
                        if hasattr(proj_layer, "dense"):
                            proj_layer = proj_layer.dense
                        if hasattr(proj_layer, "weight"):
                            w_o = proj_layer.weight.detach()
                            b_o = (
                                proj_layer.bias.detach()
                                if proj_layer.bias is not None
                                else None
                            )
                            break

        # Pattern 5: UniversalAdapter fallback
        if w_q is None or w_k is None or w_v is None or w_o is None:
            try:
                adapter = UniversalAdapter()
                if w_q is None:
                    w_q, b_q = adapter.get_weight_bias(src, "q")
                if w_k is None:
                    w_k, b_k = adapter.get_weight_bias(src, "k")
                if w_v is None:
                    w_v, b_v = adapter.get_weight_bias(src, "v")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(src, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # Strict validation
        missing = []
        if w_q is None:
            missing.append("Q projection")
        if w_k is None:
            missing.append("K projection")
        if w_v is None:
            missing.append("V projection")
        if w_o is None and not self.is_nested_encoder:
            missing.append("Out projection")

        if missing:
            raise RuntimeError(
                f"UnifiedRoPE: Failed to extract weights: {', '.join(missing)} from {type(old_attn)}"
            )

        def _copy_layer(target_layer, w, b, src_sub=None):
            if w is not None:
                if w.shape == target_layer.weight.shape:
                    target_layer.weight.copy_(w.to(target_layer.weight.dtype))
                elif w.t().contiguous().shape == target_layer.weight.shape:
                    target_layer.weight.copy_(w.t().contiguous().to(target_layer.weight.dtype))
                else:
                    min_0 = min(target_layer.weight.shape[0], w.shape[0])
                    min_1 = min(target_layer.weight.shape[1], w.shape[1])
                    target_layer.weight.zero_()
                    target_layer.weight[:min_0, :min_1].copy_(
                        w[:min_0, :min_1].to(target_layer.weight.dtype)
                    )
            if b is not None and target_layer.bias is not None:
                if b.shape == target_layer.bias.shape:
                    target_layer.bias.copy_(b.to(target_layer.bias.dtype))
                else:
                    min_b = min(target_layer.bias.shape[0], b.shape[0])
                    target_layer.bias.zero_()
                    target_layer.bias[:min_b].copy_(b[:min_b].to(target_layer.bias.dtype))
            elif target_layer.bias is not None:
                target_layer.bias.zero_()

            if (
                src_sub is not None
                and hasattr(src_sub, "lora_A")
                and hasattr(src_sub, "lora_B")
            ):
                if (
                    src_sub.lora_A is not None
                    and src_sub.lora_B is not None
                    and target_layer.lora_A is not None
                    and target_layer.lora_B is not None
                ):
                    src_lA = getattr(src_sub.lora_A, "weight", src_sub.lora_A)
                    src_lB = getattr(src_sub.lora_B, "weight", src_sub.lora_B)
                    if hasattr(src_lA, "data") and hasattr(src_lB, "data"):
                        r_min = min(target_layer.lora_A.shape[0], src_lA.shape[0])
                        in_min = min(target_layer.lora_A.shape[1], src_lA.shape[1])
                        target_layer.lora_A.data[:r_min, :in_min].copy_(
                            src_lA.data[:r_min, :in_min]
                        )
                        target_layer.lora_A.requires_grad = True

                        out_min = min(target_layer.lora_B.shape[0], src_lB.shape[0])
                        r_min_b = min(target_layer.lora_B.shape[1], src_lB.shape[1])
                        target_layer.lora_B.data[:out_min, :r_min_b].copy_(
                            src_lB.data[:out_min, :r_min_b]
                        )
                        target_layer.lora_B.requires_grad = True
                    if hasattr(src_sub, "lora_scaling"):
                        target_layer.scaling = src_sub.lora_scaling

        sub_q = getattr(src, "q_proj", getattr(src, "q_lin", getattr(src, "query", None)))
        sub_k = getattr(src, "k_proj", getattr(src, "k_lin", getattr(src, "key", None)))
        sub_v = getattr(src, "v_proj", getattr(src, "v_lin", getattr(src, "value", None)))
        sub_o = getattr(
            src,
            "o_proj",
            getattr(
                src,
                "out_lin",
                getattr(src, "dense", getattr(src, "out_proj", None)),
            ),
        )

        _copy_layer(self.q_proj, w_q, b_q, sub_q)
        _copy_layer(self.k_proj, w_k, b_k, sub_k)
        _copy_layer(self.v_proj, w_v, b_v, sub_v)
        if w_o is not None:
            _copy_layer(self.out_proj, w_o, b_o, sub_o)

            f_norm_o = float(torch.linalg.norm(self.out_proj.weight))
            if f_norm_o < 1e-4:
                logger.warning(f"UnifiedRoPE: Out projection norm is near zero: {f_norm_o}")

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

def forward(self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
head_mask: torch.Tensor | None = None,
encoder_hidden_states: torch.Tensor | None = None,
encoder_attention_mask: torch.Tensor | None = None,
past_key_value: Any | None = None,
output_attentions: bool | None = False,
output_hidden_states: bool | None = False,
return_dict: bool | None = True,
*args: Any,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor

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,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
    """Forward pass computing RoPE attention, returning (output, attention_weights)."""
    B, L, _ = hidden_states.shape

    q = (
        self.q_proj(hidden_states)
        .view(B, L, self.num_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    k = (
        self.k_proj(hidden_states)
        .view(B, L, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    v = (
        self.v_proj(hidden_states)
        .view(B, L, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # Apply RoPE
    pos_emb = kwargs.get("position_embeddings")
    if pos_emb is not None and isinstance(pos_emb, tuple) and len(pos_emb) == 2:
        cos, sin = pos_emb
        q, k = apply_rotary_pos_emb(q, k, cos, sin)
    else:
        q = self._apply_rope(q, L)
        k = self._apply_rope(k, L)

    # GQA expansion
    if self.num_kv_heads != self.num_heads:
        k = repeat_kv(k, self.num_key_value_groups)
        v = repeat_kv(v, self.num_key_value_groups)

    scale = 1.0 / math.sqrt(self.head_dim)
    scores = torch.matmul(q.float(), k.float().transpose(-1, -2).contiguous()) * scale

    # Causal mask for decoders
    if self.is_decoder and L > 1:
        causal_mask = torch.triu(
            torch.full((L, L), float("-inf"), device=scores.device, dtype=scores.dtype),
            diagonal=1,
        )
        scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

    mask = attention_mask if attention_mask is not None else kwargs.get("attn_mask")
    if mask is not None:
        scores = apply_attention_mask(scores, mask)

    attn_weights = F.softmax(scores, dim=-1)
    attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
    attn_weights = self.dropout(attn_weights.to(v.dtype))

    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:
            head_mask = head_mask
        attn_weights = attn_weights * head_mask.to(attn_weights.dtype)

    context = torch.matmul(attn_weights, v)
    context = context.transpose(1, 2).contiguous().view(B, L, self.num_heads * self.head_dim)

    if self.is_nested_encoder:
        if context.shape[-1] != self.hidden_size:
            output = self.out_proj(context)
        else:
            output = context
    else:
        output = self.out_proj(context)

    if output.shape[-1] != hidden_states.shape[-1]:
        output = (
            F.pad(output, (0, hidden_states.shape[-1] - output.shape[-1]))
            if output.shape[-1] < hidden_states.shape[-1]
            else output[..., : hidden_states.shape[-1]]
        )

    return (
        output.to(hidden_states.dtype),
        (attn_weights if output_attentions else None),
    )

Forward pass computing RoPE attention, returning (output, attention_weights).

Inherited members

• BaseAttention:
  • extra_repr

class UniversalAdapter

Expand source code

class UniversalAdapter:
    """Universal adapter supporting BERT, DistilBERT, RoBERTa, LLaMA, GPT2, DeBERTa, etc."""

    @staticmethod
    def get_weight_bias(old_attn: Any, proj: str) -> tuple:
        """Extract (weight, bias) for projection 'q', 'k', 'v', or 'o'.

        Handles:
        - RoBERTa/BERT: query/key/value attributes on .self sub-module
        - LLaMA/Qwen: q_proj/k_proj/v_proj on attention module directly
        - GPT-2: c_attn (fused, already handled by fused-QKV path); c_proj
        - DistilBERT: q_lin/k_lin/v_lin/out_lin
        - LoRALinear: unwraps .base_layer to get the frozen base weight
        - Conv1D: transposes weight from [in, out] to [out, in]
        """
        target = getattr(old_attn, "self", old_attn)
        candidate_names = CANONICAL_PROJECTION_ALIASES.get(proj, [])

        def _unwrap(mod: Any) -> Any:
            """Unwrap LoRALinear / merged adapter layers to their base_layer."""
            if mod is None:
                return mod
            # LoRALinear stores the original frozen layer as base_layer
            if hasattr(mod, "base_layer") and hasattr(mod.base_layer, "weight"):
                return mod.base_layer
            return mod

        def _extract(mod: Any):
            """Return (weight_data, bias) from a module, or None if not a valid weight-holder."""
            mod = _unwrap(mod)
            if mod is None:
                return None
            if not hasattr(mod, "weight") or mod.weight is None:
                return None
            import torch
            import torch.nn as nn

            w = mod.weight
            # Must be either an nn.Parameter or a plain Tensor (Conv1D)
            if not (isinstance(w, nn.Parameter) or isinstance(w, torch.Tensor)):
                return None
            w = w.data if hasattr(w, "data") else w
            # Projections must strictly be 2D weight matrices (excludes 1D norms, 3D convolutions, etc.)
            if w.ndim != 2:
                return None
            # Conv1D stores weights transposed: [in_features, out_features]
            is_conv1d = (mod.__class__.__name__ == "Conv1D") or hasattr(mod, "nf")
            if is_conv1d or (w.shape[0] < w.shape[1]):
                w = w.t().contiguous()
            bias = getattr(mod, "bias", None)
            bias = bias.data if (bias is not None and hasattr(bias, "data")) else bias
            # Qwen3.5 Gated Q: query is fused with gate (dim 0 has 2 * num_heads * head_dim)
            if proj == "q" and (
                hasattr(old_attn, "q_norm")
                or old_attn.__class__.__name__ == "Qwen3_5Attention"
                or hasattr(old_attn, "k_norm")
            ):
                cfg = getattr(old_attn, "config", None)
                tc = getattr(cfg, "text_config", cfg)
                h_dim = getattr(tc, "head_dim", 256)
                n_heads = getattr(tc, "num_attention_heads", getattr(old_attn, "num_heads", 8))
                if w.shape[0] == n_heads * h_dim * 2:
                    w = (
                        w.view(n_heads, 2, h_dim, -1)[:, 0, :, :]
                        .reshape(n_heads * h_dim, -1)
                        .contiguous()
                    )
                    if bias is not None:
                        bias = (
                            bias.view(n_heads, 2, h_dim)[:, 0, :]
                            .reshape(n_heads * h_dim)
                            .contiguous()
                        )
            return w, bias

        # 1. Search in target (old_attn.self or old_attn)
        for name in candidate_names:
            if hasattr(target, name):
                result = _extract(getattr(target, name))
                if result is not None:
                    return result

        # 2. Search in old_attn.output (BERT / RoBERTa output projection 'o')
        if proj == "o" and hasattr(old_attn, "output"):
            out_mod = old_attn.output
            if hasattr(out_mod, "dense"):
                result = _extract(out_mod.dense)
                if result is not None:
                    return result
            result = _extract(out_mod)
            if result is not None:
                return result

        # 3. Search directly in old_attn (catches cases where there's no .self sub-module)
        for name in candidate_names:
            if hasattr(old_attn, name):
                result = _extract(getattr(old_attn, name))
                if result is not None:
                    return result

        # 4. Fallback for Fused GPT-2 c_attn (must precede positional fallback)
        if proj in ("q", "k", "v") and hasattr(old_attn, "c_attn"):
            fused_result = _extract(old_attn.c_attn)
            if fused_result is not None:
                w, b = fused_result
                split_w = w.size(0) // 3
                split_b = b.size(0) // 3 if b is not None else 0
                if proj == "q":
                    return w[:split_w, :], (b[:split_b] if b is not None else None)
                elif proj == "k":
                    return w[split_w : 2 * split_w, :], (
                        b[split_b : 2 * split_b] if b is not None else None
                    )
                elif proj == "v":
                    return w[2 * split_w :, :], (b[2 * split_b :] if b is not None else None)

        # 4b. Fallback for Qwen3.5 / fused in_proj_qkv
        if proj in ("q", "k", "v") and hasattr(old_attn, "in_proj_qkv"):
            fused_result = _extract(old_attn.in_proj_qkv)
            if fused_result is not None:
                w, b = fused_result
                k_dim = getattr(old_attn, "key_dim", getattr(old_attn, "k_dim", w.size(0) // 3))
                v_dim = getattr(old_attn, "value_dim", getattr(old_attn, "v_dim", w.size(0) // 3))
                q_dim = w.size(0) - k_dim - v_dim if (k_dim + v_dim < w.size(0)) else k_dim
                if proj == "q":
                    return w[:q_dim, :], (b[:q_dim] if b is not None else None)
                elif proj == "k":
                    return w[q_dim : q_dim + k_dim, :], (
                        b[q_dim : q_dim + k_dim] if b is not None else None
                    )
                elif proj == "v":
                    return w[q_dim + k_dim : q_dim + k_dim + v_dim, :], (
                        b[q_dim + k_dim : q_dim + k_dim + v_dim] if b is not None else None
                    )

        # 4c. PyTorch nn.MultiheadAttention in_proj_weight / separate proj_weight
        for obj in (target, old_attn):
            if (
                proj in ("q", "k", "v")
                and hasattr(obj, "in_proj_weight")
                and obj.in_proj_weight is not None
            ):
                w = (
                    obj.in_proj_weight.data
                    if hasattr(obj.in_proj_weight, "data")
                    else obj.in_proj_weight
                )
                b = getattr(obj, "in_proj_bias", None)
                b = b.data if (b is not None and hasattr(b, "data")) else b
                split_w = w.size(0) // 3
                split_b = b.size(0) // 3 if b is not None else 0
                if proj == "q":
                    return w[:split_w, :].clone(), (b[:split_b].clone() if b is not None else None)
                elif proj == "k":
                    return w[split_w : 2 * split_w, :].clone(), (
                        b[split_b : 2 * split_b].clone() if b is not None else None
                    )
                elif proj == "v":
                    return w[2 * split_w :, :].clone(), (
                        b[2 * split_b :].clone() if b is not None else None
                    )

            if proj in ("q", "k", "v"):
                attr_w = f"{proj}_proj_weight"
                if hasattr(obj, attr_w) and getattr(obj, attr_w) is not None:
                    w = getattr(obj, attr_w)
                    w = w.data if hasattr(w, "data") else w
                    b = getattr(obj, f"{proj}_proj_bias", None)
                    b = b.data if (b is not None and hasattr(b, "data")) else b
                    return w.clone(), (b.clone() if b is not None else None)

        # 5. Positional fallback: scan all Linear children of target in declaration order.
        #    Only valid if at least 3 children exist (the typical pattern for BERT/RoBERTa variants).
        ordered_proj_map = {"q": 0, "k": 1, "v": 2}
        if proj in ordered_proj_map:
            children_with_weight = [
                (n, m) for n, m in target.named_children() if _extract(m) is not None
            ]
            if len(children_with_weight) >= 3:
                idx = ordered_proj_map[proj]
                if idx < len(children_with_weight):
                    return _extract(children_with_weight[idx][1])

        # 6. Fallback for 'o' projection when replacing leaf self-attention without local out_proj
        # (e.g. RoBERTa/BERT RobertaSelfAttention where output.dense resides in the parent container).
        # Returning an Identity matrix ensures:
        #   (a) custom modules don't crash during weight copying, and
        #   (b) context passes through transparently to parent.output.dense without double projection!
        if proj == "o":
            for mod in target.modules():
                if (
                    hasattr(mod, "weight")
                    and isinstance(getattr(mod, "weight", None), torch.Tensor)
                    and mod.weight.ndim == 2
                ):
                    d = mod.weight.shape[0]
                    eye_w = torch.eye(d, dtype=mod.weight.dtype, device=mod.weight.device)
                    zero_b = torch.zeros(d, dtype=mod.weight.dtype, device=mod.weight.device)
                    return eye_w, zero_b

        raise ValueError(
            f"Could not extract '{proj}' projection weights from {type(old_attn).__name__}"
        )

    @staticmethod
    def get_rotary_emb(old_attn: Any) -> Any | None:
        """Return rotary embeddings if present."""
        return getattr(old_attn, "rotary_emb", None)

    @staticmethod
    def get_kv_heads(config: Any) -> int:
        """Return number of key-value heads for GQA/MQA."""
        return getattr(config, "num_key_value_heads", getattr(config, "num_attention_heads", 12))

Universal adapter supporting BERT, DistilBERT, RoBERTa, LLaMA, GPT2, DeBERTa, etc.

Static methods

def get_kv_heads(config: Any) ‑> int

Expand source code

@staticmethod
def get_kv_heads(config: Any) -> int:
    """Return number of key-value heads for GQA/MQA."""
    return getattr(config, "num_key_value_heads", getattr(config, "num_attention_heads", 12))

Return number of key-value heads for GQA/MQA.

def get_rotary_emb(old_attn: Any) ‑> Any | None

Expand source code

@staticmethod
def get_rotary_emb(old_attn: Any) -> Any | None:
    """Return rotary embeddings if present."""
    return getattr(old_attn, "rotary_emb", None)

Return rotary embeddings if present.

def get_weight_bias(old_attn: Any,
proj: str) ‑> tuple

Expand source code

@staticmethod
def get_weight_bias(old_attn: Any, proj: str) -> tuple:
    """Extract (weight, bias) for projection 'q', 'k', 'v', or 'o'.

    Handles:
    - RoBERTa/BERT: query/key/value attributes on .self sub-module
    - LLaMA/Qwen: q_proj/k_proj/v_proj on attention module directly
    - GPT-2: c_attn (fused, already handled by fused-QKV path); c_proj
    - DistilBERT: q_lin/k_lin/v_lin/out_lin
    - LoRALinear: unwraps .base_layer to get the frozen base weight
    - Conv1D: transposes weight from [in, out] to [out, in]
    """
    target = getattr(old_attn, "self", old_attn)
    candidate_names = CANONICAL_PROJECTION_ALIASES.get(proj, [])

    def _unwrap(mod: Any) -> Any:
        """Unwrap LoRALinear / merged adapter layers to their base_layer."""
        if mod is None:
            return mod
        # LoRALinear stores the original frozen layer as base_layer
        if hasattr(mod, "base_layer") and hasattr(mod.base_layer, "weight"):
            return mod.base_layer
        return mod

    def _extract(mod: Any):
        """Return (weight_data, bias) from a module, or None if not a valid weight-holder."""
        mod = _unwrap(mod)
        if mod is None:
            return None
        if not hasattr(mod, "weight") or mod.weight is None:
            return None
        import torch
        import torch.nn as nn

        w = mod.weight
        # Must be either an nn.Parameter or a plain Tensor (Conv1D)
        if not (isinstance(w, nn.Parameter) or isinstance(w, torch.Tensor)):
            return None
        w = w.data if hasattr(w, "data") else w
        # Projections must strictly be 2D weight matrices (excludes 1D norms, 3D convolutions, etc.)
        if w.ndim != 2:
            return None
        # Conv1D stores weights transposed: [in_features, out_features]
        is_conv1d = (mod.__class__.__name__ == "Conv1D") or hasattr(mod, "nf")
        if is_conv1d or (w.shape[0] < w.shape[1]):
            w = w.t().contiguous()
        bias = getattr(mod, "bias", None)
        bias = bias.data if (bias is not None and hasattr(bias, "data")) else bias
        # Qwen3.5 Gated Q: query is fused with gate (dim 0 has 2 * num_heads * head_dim)
        if proj == "q" and (
            hasattr(old_attn, "q_norm")
            or old_attn.__class__.__name__ == "Qwen3_5Attention"
            or hasattr(old_attn, "k_norm")
        ):
            cfg = getattr(old_attn, "config", None)
            tc = getattr(cfg, "text_config", cfg)
            h_dim = getattr(tc, "head_dim", 256)
            n_heads = getattr(tc, "num_attention_heads", getattr(old_attn, "num_heads", 8))
            if w.shape[0] == n_heads * h_dim * 2:
                w = (
                    w.view(n_heads, 2, h_dim, -1)[:, 0, :, :]
                    .reshape(n_heads * h_dim, -1)
                    .contiguous()
                )
                if bias is not None:
                    bias = (
                        bias.view(n_heads, 2, h_dim)[:, 0, :]
                        .reshape(n_heads * h_dim)
                        .contiguous()
                    )
        return w, bias

    # 1. Search in target (old_attn.self or old_attn)
    for name in candidate_names:
        if hasattr(target, name):
            result = _extract(getattr(target, name))
            if result is not None:
                return result

    # 2. Search in old_attn.output (BERT / RoBERTa output projection 'o')
    if proj == "o" and hasattr(old_attn, "output"):
        out_mod = old_attn.output
        if hasattr(out_mod, "dense"):
            result = _extract(out_mod.dense)
            if result is not None:
                return result
        result = _extract(out_mod)
        if result is not None:
            return result

    # 3. Search directly in old_attn (catches cases where there's no .self sub-module)
    for name in candidate_names:
        if hasattr(old_attn, name):
            result = _extract(getattr(old_attn, name))
            if result is not None:
                return result

    # 4. Fallback for Fused GPT-2 c_attn (must precede positional fallback)
    if proj in ("q", "k", "v") and hasattr(old_attn, "c_attn"):
        fused_result = _extract(old_attn.c_attn)
        if fused_result is not None:
            w, b = fused_result
            split_w = w.size(0) // 3
            split_b = b.size(0) // 3 if b is not None else 0
            if proj == "q":
                return w[:split_w, :], (b[:split_b] if b is not None else None)
            elif proj == "k":
                return w[split_w : 2 * split_w, :], (
                    b[split_b : 2 * split_b] if b is not None else None
                )
            elif proj == "v":
                return w[2 * split_w :, :], (b[2 * split_b :] if b is not None else None)

    # 4b. Fallback for Qwen3.5 / fused in_proj_qkv
    if proj in ("q", "k", "v") and hasattr(old_attn, "in_proj_qkv"):
        fused_result = _extract(old_attn.in_proj_qkv)
        if fused_result is not None:
            w, b = fused_result
            k_dim = getattr(old_attn, "key_dim", getattr(old_attn, "k_dim", w.size(0) // 3))
            v_dim = getattr(old_attn, "value_dim", getattr(old_attn, "v_dim", w.size(0) // 3))
            q_dim = w.size(0) - k_dim - v_dim if (k_dim + v_dim < w.size(0)) else k_dim
            if proj == "q":
                return w[:q_dim, :], (b[:q_dim] if b is not None else None)
            elif proj == "k":
                return w[q_dim : q_dim + k_dim, :], (
                    b[q_dim : q_dim + k_dim] if b is not None else None
                )
            elif proj == "v":
                return w[q_dim + k_dim : q_dim + k_dim + v_dim, :], (
                    b[q_dim + k_dim : q_dim + k_dim + v_dim] if b is not None else None
                )

    # 4c. PyTorch nn.MultiheadAttention in_proj_weight / separate proj_weight
    for obj in (target, old_attn):
        if (
            proj in ("q", "k", "v")
            and hasattr(obj, "in_proj_weight")
            and obj.in_proj_weight is not None
        ):
            w = (
                obj.in_proj_weight.data
                if hasattr(obj.in_proj_weight, "data")
                else obj.in_proj_weight
            )
            b = getattr(obj, "in_proj_bias", None)
            b = b.data if (b is not None and hasattr(b, "data")) else b
            split_w = w.size(0) // 3
            split_b = b.size(0) // 3 if b is not None else 0
            if proj == "q":
                return w[:split_w, :].clone(), (b[:split_b].clone() if b is not None else None)
            elif proj == "k":
                return w[split_w : 2 * split_w, :].clone(), (
                    b[split_b : 2 * split_b].clone() if b is not None else None
                )
            elif proj == "v":
                return w[2 * split_w :, :].clone(), (
                    b[2 * split_b :].clone() if b is not None else None
                )

        if proj in ("q", "k", "v"):
            attr_w = f"{proj}_proj_weight"
            if hasattr(obj, attr_w) and getattr(obj, attr_w) is not None:
                w = getattr(obj, attr_w)
                w = w.data if hasattr(w, "data") else w
                b = getattr(obj, f"{proj}_proj_bias", None)
                b = b.data if (b is not None and hasattr(b, "data")) else b
                return w.clone(), (b.clone() if b is not None else None)

    # 5. Positional fallback: scan all Linear children of target in declaration order.
    #    Only valid if at least 3 children exist (the typical pattern for BERT/RoBERTa variants).
    ordered_proj_map = {"q": 0, "k": 1, "v": 2}
    if proj in ordered_proj_map:
        children_with_weight = [
            (n, m) for n, m in target.named_children() if _extract(m) is not None
        ]
        if len(children_with_weight) >= 3:
            idx = ordered_proj_map[proj]
            if idx < len(children_with_weight):
                return _extract(children_with_weight[idx][1])

    # 6. Fallback for 'o' projection when replacing leaf self-attention without local out_proj
    # (e.g. RoBERTa/BERT RobertaSelfAttention where output.dense resides in the parent container).
    # Returning an Identity matrix ensures:
    #   (a) custom modules don't crash during weight copying, and
    #   (b) context passes through transparently to parent.output.dense without double projection!
    if proj == "o":
        for mod in target.modules():
            if (
                hasattr(mod, "weight")
                and isinstance(getattr(mod, "weight", None), torch.Tensor)
                and mod.weight.ndim == 2
            ):
                d = mod.weight.shape[0]
                eye_w = torch.eye(d, dtype=mod.weight.dtype, device=mod.weight.device)
                zero_b = torch.zeros(d, dtype=mod.weight.dtype, device=mod.weight.device)
                return eye_w, zero_b

    raise ValueError(
        f"Could not extract '{proj}' projection weights from {type(old_attn).__name__}"
    )

Extract (weight, bias) for projection 'q', 'k', 'v', or 'o'.

Handles: - RoBERTa/BERT: query/key/value attributes on .self sub-module - LLaMA/Qwen: q_proj/k_proj/v_proj on attention module directly - GPT-2: c_attn (fused, already handled by fused-QKV path); c_proj - DistilBERT: q_lin/k_lin/v_lin/out_lin - LoRALinear: unwraps .base_layer to get the frozen base weight - Conv1D: transposes weight from [in, out] to [out, in]

class VSpecAttention (config: Any | None = None,
rank: int = 64,
pretrained_attn: nn.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "vspec", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "vspec_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class VSpecAttention(BaseAttention):
    r"""VSpec++ Attention: Variational Spectral Token-Adaptive Attention Module."""

    supports_gqa: bool = True
    supports_causal: bool = True
    supports_bidirectional: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int = 64,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        src_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")

        cfg_head_dim = _cfg_get(config, "head_dim", kwargs.get("head_dim"))
        if cfg_head_dim is None and src_attn is not None:
            for attr in ("q_proj", "query", "q_lin"):
                mod = getattr(src_attn, attr, None)
                if mod is not None and hasattr(mod, "weight"):
                    num_h = _cfg_get(config, "num_attention_heads", kwargs.get("num_heads", 12))
                    cfg_head_dim = mod.weight.shape[0] // max(1, num_h)
                    break

        _mt = str(getattr(config, "model_type", "") or "").lower() if config else ""
        _decoder_prefixes = (
            "gpt",
            "qwen",
            "llama",
            "mistral",
            "gemma",
            "phi",
            "opt",
            "bloom",
            "falcon",
            "olmo",
            "mpt",
            "stablelm",
        )
        is_dec = (
            (
                getattr(config, "is_decoder", False)
                or getattr(config, "causal", False)
                or any(_mt.startswith(p) for p in _decoder_prefixes)
            )
            if config
            else False
        )
        if hasattr(config, "architectures") and any(
            "GPT" in str(a) or "CausalLM" in str(a)
            for a in (getattr(config, "architectures", []) or [])
        ):
            is_dec = True
        if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False):
            is_dec = True

        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 = (
                int(cfg_head_dim) if cfg_head_dim is not None else 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 = is_dec
        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 = (
                int(cfg_head_dim)
                if cfg_head_dim is not None
                else (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 = is_dec

        # Explicit overrides
        if "hidden_size" in kwargs:
            self.hidden_size = kwargs["hidden_size"]
        if "num_heads" in kwargs:
            self.num_heads = kwargs["num_heads"]
        if "head_dim" in kwargs:
            self.head_dim = kwargs["head_dim"]
        elif cfg_head_dim is not None:
            self.head_dim = int(cfg_head_dim)
        else:
            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))

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)
        self.is_causal = self.is_decoder

        self.is_nested_encoder: bool = False
        self.rank: int = rank
        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))

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and src_attn is not None:
            target_dtype = next(src_attn.parameters(), torch.tensor([])).dtype
            if not isinstance(target_dtype, torch.dtype):
                target_dtype = torch.float32

        self.Pq = DynamicLowRankProjector(
            rank=self.rank,
            in_dim=self.hidden_size,
            out_dim=self.all_head_size,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.Pk = DynamicLowRankProjector(
            rank=self.rank,
            in_dim=self.hidden_size,
            out_dim=self.kv_head_size,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.Pv = DynamicLowRankProjector(
            rank=self.rank,
            in_dim=self.hidden_size,
            out_dim=self.kv_head_size,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.Po = DynamicLowRankProjector(
            rank=self.rank,
            in_dim=self.all_head_size,
            out_dim=self.hidden_size,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
            dtype=target_dtype,
        )

        self.sg_q = VariationalSpectralGate(self.hidden_size, self.rank)
        self.sg_k = VariationalSpectralGate(self.hidden_size, self.rank)
        self.sg_v = VariationalSpectralGate(self.hidden_size, self.rank)

        self.head_gate = HeadGate(self.hidden_size, self.num_heads)
        self.router = TokenRouter(self.hidden_size, keep_ratio=0.75)

        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)

        if target_dtype is not None:
            self.to(target_dtype)

    @property
    def query(self):
        return self.Pq

    @property
    def key(self):
        return self.Pk

    @property
    def value(self):
        return self.Pv

    @property
    def q_proj(self):
        return self.Pq

    @property
    def k_proj(self):
        return self.Pk

    @property
    def v_proj(self):
        return self.Pv

    @property
    def out_proj(self):
        return self.Po

    @property
    def o_proj(self):
        return self.Po

    @property
    def dense(self):
        return self.Po

    @property
    def c_proj(self):
        return self.Po

    def copy_weights(self, src: Any) -> None:
        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)
                        base = getattr(layer, "base_layer", layer)
                        if hasattr(base, "weight"):
                            extracted[proj] = (
                                base.weight.data,
                                getattr(base, "bias", None),
                            )
                            break

        # Detect nested encoder architecture
        cls_name = 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(
                "VSpecAttention missing required weights for projections: %s; keeping defaults",
                missing,
            )

        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

        # Copy into projectors
        if extracted["q"][0] is not None:
            self.Pq.set_weight(_orient(extracted["q"][0], self.hidden_size), extracted["q"][1])
        if extracted["k"][0] is not None:
            self.Pk.set_weight(_orient(extracted["k"][0], self.hidden_size), extracted["k"][1])
        if extracted["v"][0] is not None:
            self.Pv.set_weight(_orient(extracted["v"][0], self.hidden_size), extracted["v"][1])

        if extracted["o"][0] is not None and not self.is_nested_encoder:
            self.Po.set_weight(_orient(extracted["o"][0], self.all_head_size), extracted["o"][1])
        elif self.is_nested_encoder:
            eye_w = torch.eye(self.all_head_size, self.all_head_size)
            self.Po.set_weight(eye_w, None)

        # Preserve source LoRA adapters if present
        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.Pq),
            ("k", self.Pk),
            ("v", self.Pv),
            ("o", self.Po),
        ]:
            for name in source_submodules[proj_key]:
                if hasattr(src, name):
                    sub = getattr(src, name)
                    target.copy_lora_from(sub)
                    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] | torch.Tensor:
        if encoder_hidden_states is not None or kwargs.get("encoder_hidden_states") is not None:
            raise NotImplementedError("Cross-attention is not implemented for VSpecAttention")

        B, N, D = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        # Mask handling
        mask = attention_mask if attention_mask is not None else kwargs.get("mask")
        valid = get_token_valid_mask(mask, B, N, device) if mask is not None else None

        is_decoder = (
            self.is_decoder
            or getattr(self.config, "is_decoder", False)
            or getattr(self.config, "causal", False)
            or any(
                str(getattr(self.config, "model_type", "") or "").lower().startswith(p)
                for p in (
                    "gpt",
                    "qwen",
                    "llama",
                    "mistral",
                    "gemma",
                    "phi",
                    "opt",
                    "bloom",
                    "falcon",
                    "olmo",
                    "mpt",
                    "stablelm",
                )
            )
        )

        # Token routing: Skip for nested encoders (RoBERTa/BERT) to preserve context
        if self.is_nested_encoder:
            routed = hidden_states
            idx = torch.arange(N, device=device).unsqueeze(0).expand(B, -1)
            k_len = N
            valid_routed = valid
        else:
            routed, idx = self.router(hidden_states, valid, is_decoder=is_decoder)
            k_len = routed.shape[1]
            valid_routed = torch.gather(valid, 1, idx) if valid is not None else None

        # Variational spectral gates using masked mean over routed tokens
        gq, _, _ = self.sg_q(routed, valid_routed)
        gk, _, _ = self.sg_k(routed, valid_routed)
        gv, _, _ = self.sg_v(routed, valid_routed)

        Q = self.Pq(routed, gq)
        K = self.Pk(routed, gk)
        V = self.Pv(routed, gv)

        H = self.num_heads
        H_kv = self.num_kv_heads
        Dh = self.head_dim

        Q = Q.view(B, k_len, H, Dh).transpose(1, 2).contiguous()
        K = K.view(B, k_len, H_kv, Dh).transpose(1, 2).contiguous()
        V = V.view(B, k_len, H_kv, Dh).transpose(1, 2).contiguous()

        # Rotary position embeddings
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
            cos_b = cos.expand(B, -1, -1) if cos.size(0) == 1 else cos
            sin_b = sin.expand(B, -1, -1) if sin.size(0) == 1 else sin
            cos_routed = torch.gather(cos_b, 1, idx.unsqueeze(-1).expand(-1, -1, cos.size(-1)))
            sin_routed = torch.gather(sin_b, 1, idx.unsqueeze(-1).expand(-1, -1, sin.size(-1)))
            Q, K = apply_rotary_pos_emb(Q, K, cos_routed, sin_routed)
        elif getattr(self, "rotary_emb", None) is not None:
            try:
                from neuralops.utils.compat import maybe_apply_rotary

                Q, K = maybe_apply_rotary(
                    Q,
                    K,
                    getattr(self, "adapter", None),
                    self.rotary_emb,
                    kwargs.get("position_ids"),
                )
            except Exception as e:
                logger.warning(f"maybe_apply_rotary failed: {e}")

        # GQA expansion
        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        hg = self.head_gate(routed, valid_routed).unsqueeze(-1).unsqueeze(-1)

        scale = 1.0 / math.sqrt(Dh)
        scores = (
            torch.matmul(Q.float(), K.float().transpose(-1, -2).contiguous()) * scale * hg.float()
        )

        # Check causal decoder
        if is_decoder and k_len > 1:
            causal_mask = torch.triu(
                torch.full((k_len, k_len), float("-inf"), device=device, dtype=scores.dtype),
                diagonal=1,
            )
            scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)

        # Apply padding mask to attention scores over routed tokens
        if valid_routed is not None:
            pad_mask = (
                (~valid_routed).unsqueeze(1).unsqueeze(2).expand(-1, self.num_heads, k_len, k_len)
            )
            scores = scores.masked_fill(pad_mask, float("-inf"))

        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = torch.nan_to_num(attn_weights, nan=0.0).to(dtype)
        attn_weights = self.dropout(attn_weights)

        # Head mask broadcasting
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.view(1, -1, 1, 1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.view(head_mask.size(0), head_mask.size(1), 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            elif head_mask.dim() == 4:
                if (
                    head_mask.shape[-1] == head_mask.shape[-2]
                    and head_mask.shape[-1] != attn_weights.shape[-1]
                ):
                    head_mask = head_mask.mean(dim=-1, keepdim=True)
            attn_weights = attn_weights * head_mask.to(attn_weights.dtype)

        context = torch.matmul(attn_weights, V)
        context = context.transpose(1, 2).contiguous().view(B, k_len, self.all_head_size)

        if self.is_nested_encoder:
            return (context, (attn_weights if output_attentions else None))

        out = self.Po(context, gq)

        if out.shape[-1] != D:
            out = F.pad(out, (0, D - out.shape[-1])) if out.shape[-1] < D else out[..., :D]

        # Scatter back to full sequence length preserving residual:
        # Non-routed tokens have zero attention delta, so outer block residual keeps them unchanged.
        if k_len == N:
            output = out
        else:
            zeros = torch.zeros(B, N, D, device=device, dtype=out.dtype)
            output = zeros.scatter(1, idx.unsqueeze(-1).expand(-1, -1, D), out)

        return (output, (attn_weights if output_attentions else None))

VSpec++ Attention: Variational Spectral Token-Adaptive Attention Module.

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.Po

prop dense

Expand source code

@property
def dense(self):
    return self.Po

prop k_proj

Expand source code

@property
def k_proj(self):
    return self.Pk

prop key

Expand source code

@property
def key(self):
    return self.Pk

prop o_proj

Expand source code

@property
def o_proj(self):
    return self.Po

prop out_proj

Expand source code

@property
def out_proj(self):
    return self.Po

prop q_proj

Expand source code

@property
def q_proj(self):
    return self.Pq

prop query

Expand source code

@property
def query(self):
    return self.Pq

prop v_proj

Expand source code

@property
def v_proj(self):
    return self.Pv

prop value

Expand source code

@property
def value(self):
    return self.Pv

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class VariationalSpectralGate (hidden_dim: int, rank: int)

Expand source code

class VariationalSpectralGate(nn.Module):
    r"""Input-dependent variational gate over truncated singular values with masked pooling."""

    def __init__(self, hidden_dim: int, rank: int) -> None:
        super().__init__()
        self.net = nn.Sequential(nn.Linear(hidden_dim, 64), nn.GELU(), nn.Linear(64, rank * 2))
        # Initialize gate to neutral (near 1.0) so SVD representations are preserved at step 0
        nn.init.zeros_(self.net[-1].weight)
        nn.init.constant_(self.net[-1].bias[:rank], 3.0)  # mu -> sigmoid(3.0) ≈ 0.95
        nn.init.constant_(self.net[-1].bias[rank:], -3.0)  # log_sigma -> low variance

    def forward(
        self, x: torch.Tensor, valid_mask: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        net_dtype = next(self.net.parameters()).dtype
        pooled = _masked_mean(x, valid_mask).to(net_dtype)
        stats = self.net(pooled)
        mu, log_sigma = stats.chunk(2, dim=-1)
        log_sigma = torch.clamp(log_sigma, min=-10.0, max=2.0)
        if self.training:
            eps = torch.randn_like(mu)
            z = mu + eps * torch.exp(log_sigma)
        else:
            z = mu
        g = torch.sigmoid(z).to(x.dtype)
        return g, mu, log_sigma

Input-dependent variational gate over truncated singular values with masked pooling.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor, valid_mask: torch.Tensor | None = None) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Expand source code

def forward(
    self, x: torch.Tensor, valid_mask: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    net_dtype = next(self.net.parameters()).dtype
    pooled = _masked_mean(x, valid_mask).to(net_dtype)
    stats = self.net(pooled)
    mu, log_sigma = stats.chunk(2, dim=-1)
    log_sigma = torch.clamp(log_sigma, min=-10.0, max=2.0)
    if self.training:
        eps = torch.randn_like(mu)
        z = mu + eps * torch.exp(log_sigma)
    else:
        z = mu
    g = torch.sigmoid(z).to(x.dtype)
    return g, mu, log_sigma

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ViGAttention (config: Any = None,
rank: int | None = None,
pretrained_attn: nn.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "vig", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "vig_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class ViGAttention(BaseAttention):
    r"""Vision GNN-inspired 2D Grid Linear Attention with Learned Gated Fusion."""

    supports_gqa: bool = True
    supports_causal: bool = True
    supports_bidirectional: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any = None,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        src_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")

        if config is not None and not isinstance(config, int):
            self.adapter = ModelAdapter(config)
            self.num_heads: int = self.adapter.num_heads
            self.num_kv_heads: int = self.adapter.num_key_value_heads
            self.dim: int = self.adapter.hidden_size
            self.head_dim: int = self.adapter.head_dim
            self.all_head_size: int = self.num_heads * self.head_dim
            self.kv_head_size: int = self.num_kv_heads * self.head_dim
            self.kv_repeat: int = self.adapter.kv_repeat
            self.is_decoder: bool = getattr(config, "is_decoder", False)
        else:
            self.adapter = None
            self.dim = _cfg_get(
                config,
                "hidden_size",
                kwargs.get("hidden_size", kwargs.get("d_model", 768)),
            )
            self.num_heads = _cfg_get(config, "num_attention_heads", kwargs.get("num_heads", 12))
            self.num_kv_heads = _cfg_get(config, "num_key_value_heads", self.num_heads)
            self.head_dim = self.dim // max(1, self.num_heads)
            self.all_head_size = self.num_heads * self.head_dim
            self.kv_head_size = self.num_kv_heads * self.head_dim
            self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
            self.is_decoder = bool(_cfg_get(config, "is_decoder", kwargs.get("is_decoder", False)))

        # Explicit overrides
        if "hidden_size" in kwargs:
            self.dim = kwargs["hidden_size"]
            self.head_dim = self.dim // max(1, self.num_heads)
            self.all_head_size = self.num_heads * self.head_dim
            self.kv_head_size = self.num_kv_heads * self.head_dim
        if "num_heads" in kwargs:
            self.num_heads = kwargs["num_heads"]
            self.head_dim = self.dim // max(1, self.num_heads)
            self.all_head_size = self.num_heads * self.head_dim
            self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)

        self.is_nested_encoder: bool = False
        self.lora_rank = kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank))
        self.lora_alpha = kwargs.get("lora_alpha", _cfg_get(config, "lora_alpha", lora_alpha))

        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        self.gate_net = nn.Linear(self.dim, 2 * self.num_heads)
        dropout_p = _cfg_get(
            config,
            "dropout",
            _cfg_get(
                config,
                "hidden_dropout_prob",
                _cfg_get(config, "attention_probs_dropout_prob", 0.1),
            ),
        )
        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(src_attn, "rotary_emb", None)

        if src_attn is not None:
            self.copy_weights(src_attn)

    # State dict compatibility aliases via properties
    @property
    def query(self) -> LoRALinear:
        return self.q_proj

    @property
    def key(self) -> LoRALinear:
        return self.k_proj

    @property
    def value(self) -> LoRALinear:
        return self.v_proj

    @property
    def out(self) -> LoRALinear:
        return self.out_proj

    @property
    def W_q(self) -> LoRALinear:
        return self.q_proj

    @property
    def W_k(self) -> LoRALinear:
        return self.k_proj

    @property
    def W_v(self) -> LoRALinear:
        return self.v_proj

    @property
    def q_lin(self) -> LoRALinear:
        return self.q_proj

    @property
    def k_lin(self) -> LoRALinear:
        return self.k_proj

    @property
    def v_lin(self) -> LoRALinear:
        return self.v_proj

    @property
    def out_lin(self) -> LoRALinear:
        return self.out_proj

    @property
    def dense(self) -> LoRALinear:
        return self.out_proj

    @property
    def o_proj(self) -> LoRALinear:
        return self.out_proj

    @property
    def c_proj(self) -> LoRALinear:
        return self.out_proj

    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. UniversalAdapter extraction
        for proj in ["q", "k", "v", "o"]:
            try:
                w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
                if w is not None:
                    extracted[proj] = (w, b)
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # 2. GPT-2 Fused c_attn Conv1D handling
        if hasattr(old_attn, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            weight = old_attn.c_attn.weight.data
            bias = getattr(old_attn.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if (
                ("Conv1D" in old_attn.c_attn.__class__.__name__)
                or hasattr(old_attn.c_attn, "nf")
                or (
                    weight.dim() == 2
                    and (weight.shape[0] < weight.shape[1] or weight.shape[1] == 3 * self.dim)
                )
            ):
                weight = weight.t().contiguous()
            split_size = weight.shape[0] // 3
            q_w = weight[:split_size, :]
            k_w = weight[split_size : 2 * split_size, :]
            v_w = weight[2 * split_size :, :]
            q_b = bias_data[:split_size] if bias_data is not None else None
            k_b = bias_data[split_size : 2 * split_size] if bias_data is not None else None
            v_b = bias_data[2 * split_size :] if bias_data is not None else None
            extracted["q"] = (q_w, q_b)
            extracted["k"] = (k_w, k_b)
            extracted["v"] = (v_w, v_b)

        if hasattr(old_attn, "c_proj") and extracted["o"][0] is None:
            w_o = old_attn.c_proj.weight.data
            if (
                ("Conv1D" in old_attn.c_proj.__class__.__name__)
                or hasattr(old_attn.c_proj, "nf")
                or (
                    w_o.dim() == 2
                    and (
                        w_o.shape[0] < w_o.shape[1]
                        or (w_o.shape[0] != self.dim and w_o.shape[1] == self.dim)
                    )
                )
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(old_attn.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(old_attn, "self") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

        if (
            hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
            and extracted["o"][0] is None
        ):
            dense = old_attn.output.dense
            extracted["o"] = (dense.weight.data, getattr(dense, "bias", None))

        # 4. Modern Decoders (Qwen3.5, LLaMA, Mistral) & DistilBERT
        named_map = {
            "q": ["q_proj", "q_lin", "query", "W_q"],
            "k": ["k_proj", "k_lin", "key", "W_k"],
            "v": ["v_proj", "v_lin", "value", "W_v"],
            "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        base = getattr(layer, "base_layer", layer)
                        if hasattr(base, "weight"):
                            extracted[proj] = (
                                base.weight.data,
                                getattr(base, "bias", None),
                            )
                            break

        # Detect nested encoder architecture
        cls_name = old_attn.__class__.__name__
        has_out_proj = any(
            hasattr(old_attn, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
                "out",
            )
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ) and not has_out_proj

        missing = [k for k in ["q", "k", "v"] if extracted[k][0] is None]
        if not self.is_nested_encoder and extracted["o"][0] is None:
            missing.append("o")
        if missing:
            logger.warning(
                "Attention module copy_weights warning. Proceeding with initialized weights."
            )

        def _orient(w: torch.Tensor, in_dim: int) -> torch.Tensor:
            if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim:
                return w.t().contiguous()
            return w

        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.dim
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape
                    target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                        w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                    )
                    if b is not None and target.bias is not None:
                        b_val = b.data if isinstance(b, nn.Parameter) else b
                        target.bias.data[: min(len(target.bias), len(b_val))].copy_(
                            b_val[: min(len(target.bias), len(b_val))]
                        )
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                # Frobenius norm verification (>= 1e-4)
                w_norm = float(torch.linalg.norm(target.weight.data.float()).item())
                if w_norm < 1e-4:
                    logger.warning(
                        f"ViGAttention weight verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                    )

                target.freeze_base_weights()

        # Preserve source LoRA adapters
        source_submodules = {
            "q": ["q_proj", "q_lin", "query", "W_q"],
            "k": ["k_proj", "k_lin", "key", "W_k"],
            "v": ["v_proj", "v_lin", "value", "W_v"],
            "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
        }
        for proj_key, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            for name in source_submodules[proj_key]:
                if hasattr(old_attn, name):
                    sub = getattr(old_attn, name)
                    if (
                        hasattr(sub, "lora_A")
                        and hasattr(sub, "lora_B")
                        and sub.lora_A is not None
                        and sub.lora_B is not None
                    ):
                        src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                        src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                        if hasattr(src_lA, "data") and target.lora_A is not None:
                            if (
                                src_lA.shape[0] == target.lora_A.shape[0]
                                and src_lA.shape[1] == target.lora_A.shape[1]
                            ):
                                target.lora_A.data.copy_(src_lA.data)
                                target.lora_A.requires_grad = True
                            else:
                                print(
                                    f"Warning: lora_A shape mismatch for {proj_key}; keeping zero-init."
                                )
                        if hasattr(src_lB, "data") and target.lora_B is not None:
                            if (
                                src_lB.shape[0] == target.lora_B.shape[0]
                                and src_lB.shape[1] == target.lora_B.shape[1]
                            ):
                                target.lora_B.data.copy_(src_lB.data)
                                target.lora_B.requires_grad = True
                            else:
                                print(
                                    f"Warning: lora_B shape mismatch for {proj_key}; keeping zero-init."
                                )
                        if hasattr(sub, "lora_scaling"):
                            target.lora_scaling = sub.lora_scaling
                        break

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        B, N, D = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        Hh = self.num_heads
        H_kv = self.num_kv_heads
        Dh = self.head_dim

        q = self.q_proj(hidden_states).view(B, N, Hh, Dh).transpose(1, 2).contiguous()
        k = self.k_proj(hidden_states).view(B, N, H_kv, Dh).transpose(1, 2).contiguous()
        v = self.v_proj(hidden_states).view(B, N, H_kv, Dh).transpose(1, 2).contiguous()

        # Rotary position embeddings
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
            q, k = apply_rotary_pos_emb(q, k, cos, sin)
        elif getattr(self, "rotary_emb", None) is not None:
            try:
                from neuralops.utils.compat import maybe_apply_rotary

                q, k = maybe_apply_rotary(
                    q,
                    k,
                    getattr(self, "adapter", None),
                    self.rotary_emb,
                    kwargs.get("position_ids"),
                )
            except Exception as e:
                logger.warning("maybe_apply_rotary failed: %s; continuing without RoPE", e)

        # GQA expansion
        if self.kv_repeat > 1:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        is_decoder = (
            self.is_decoder
            or getattr(getattr(self, "config", None), "is_decoder", False)
            or getattr(getattr(self, "config", None), "model_type", "")
            in ("gpt2", "qwen", "qwen2", "qwen3_5", "llama")
        )

        # 1. Topological Graph Adjacency for 1D Token Sequences
        # Bounded local window (radius 16) + Global Anchor Tokens (token 0 and token N-1)
        pos = torch.arange(N, device=device)
        diff = pos.unsqueeze(0) - pos.unsqueeze(1)  # diff[i, j] = pos[j] - pos[i]
        local_mask = diff.abs() <= 16

        anchor_mask = torch.zeros((N, N), device=device, dtype=torch.bool)
        anchor_mask[0, :] = True
        anchor_mask[:, 0] = True
        if N > 1:
            anchor_mask[-1, :] = True
            anchor_mask[:, -1] = True

        graph_adj = local_mask | anchor_mask
        if is_decoder:
            causal_mask = diff <= 0  # i >= j
            graph_adj = graph_adj & causal_mask

        # Branch 1: Topological Graph Attention
        scale = 1.0 / math.sqrt(Dh)
        scores_graph = torch.matmul(q, k.transpose(-1, -2).contiguous()) * scale
        scores_graph = scores_graph.masked_fill(~graph_adj.unsqueeze(0).unsqueeze(0), -10000.0)
        if attention_mask is not None:
            scores_graph = apply_attention_mask(scores_graph, attention_mask)
        attn_graph = torch.nan_to_num(F.softmax(scores_graph.float(), dim=-1), nan=0.0).to(dtype)
        out_graph = torch.matmul(attn_graph, v)

        # Branch 2: Global Linear Recurrence / Attention with denom.clamp(min=1.0)
        q_lin = F.elu(q.float()) + 1.0
        k_lin = F.elu(k.float()) + 1.0

        valid_m = get_token_valid_mask(attention_mask, B, N, device)
        if valid_m is not None:
            k_lin = k_lin * valid_m.unsqueeze(1).unsqueeze(-1).to(k_lin.dtype)
            v_f = v.float() * valid_m.unsqueeze(1).unsqueeze(-1).to(v.dtype)
        else:
            v_f = v.float()

        if is_decoder:
            causal_tril = (
                torch.tril(torch.ones(N, N, device=device, dtype=torch.bool))
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores_lin = torch.matmul(q_lin, k_lin.transpose(-1, -2).contiguous()) * causal_tril
            denom_lin = scores_lin.sum(dim=-1, keepdim=True)
            out_global = (torch.matmul(scores_lin, v_f) / denom_lin.clamp(min=1.0)).to(dtype)
        else:
            kv = torch.matmul(k_lin.transpose(-1, -2).contiguous(), v_f)  # [B, H, Dh, Dh]
            num = torch.matmul(q_lin, kv)  # [B, H, N, Dh]
            denom = torch.matmul(
                q_lin, k_lin.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous()
            )  # [B, H, N, 1]
            out_global = (num / denom.clamp(min=1.0)).to(dtype)

        # Gated Fusion
        gates = torch.sigmoid(self.gate_net(hidden_states)).view(B, N, self.num_heads, 2)
        gate_graph = gates[..., 0].transpose(1, 2).contiguous().unsqueeze(-1).to(dtype)
        gate_global = gates[..., 1].transpose(1, 2).contiguous().unsqueeze(-1).to(dtype)

        out = gate_graph * out_graph + gate_global * out_global
        out = self.dropout(out)

        # Head mask broadcasting
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.view(1, -1, 1, 1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.view(head_mask.size(0), -1, 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            out = out * head_mask.to(dtype)

        out = out.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)

        return (output, None)

Vision GNN-inspired 2D Grid Linear Attention with Learned Gated Fusion.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• BaseAttention
• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

var supports_bidirectional : bool

The type of the None singleton.

var supports_causal : bool

The type of the None singleton.

var supports_gqa : bool

The type of the None singleton.

var supports_rope : bool

The type of the None singleton.

Instance variables

prop W_k : LoRALinear

Expand source code

@property
def W_k(self) -> LoRALinear:
    return self.k_proj

prop W_q : LoRALinear

Expand source code

@property
def W_q(self) -> LoRALinear:
    return self.q_proj

prop W_v : LoRALinear

Expand source code

@property
def W_v(self) -> LoRALinear:
    return self.v_proj

prop c_proj : LoRALinear

Expand source code

@property
def c_proj(self) -> LoRALinear:
    return self.out_proj

prop dense : LoRALinear

Expand source code

@property
def dense(self) -> LoRALinear:
    return self.out_proj

prop k_lin : LoRALinear

Expand source code

@property
def k_lin(self) -> LoRALinear:
    return self.k_proj

prop key : LoRALinear

Expand source code

@property
def key(self) -> LoRALinear:
    return self.k_proj

prop o_proj : LoRALinear

Expand source code

@property
def o_proj(self) -> LoRALinear:
    return self.out_proj

prop out : LoRALinear

Expand source code

@property
def out(self) -> LoRALinear:
    return self.out_proj

prop out_lin : LoRALinear

Expand source code

@property
def out_lin(self) -> LoRALinear:
    return self.out_proj

prop q_lin : LoRALinear

Expand source code

@property
def q_lin(self) -> LoRALinear:
    return self.q_proj

prop query : LoRALinear

Expand source code

@property
def query(self) -> LoRALinear:
    return self.q_proj

prop v_lin : LoRALinear

Expand source code

@property
def v_lin(self) -> LoRALinear:
    return self.v_proj

prop value : LoRALinear

Expand source code

@property
def value(self) -> LoRALinear:
    return self.v_proj

Methods

def copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

    # 1. UniversalAdapter extraction
    for proj in ["q", "k", "v", "o"]:
        try:
            w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
            if w is not None:
                extracted[proj] = (w, b)
        except Exception as e:
            logger.debug("UniversalAdapter fallback failed: %s", e)

    # 2. GPT-2 Fused c_attn Conv1D handling
    if hasattr(old_attn, "c_attn") and (
        extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
    ):
        weight = old_attn.c_attn.weight.data
        bias = getattr(old_attn.c_attn, "bias", None)
        bias_data = bias.data if bias is not None else None
        if (
            ("Conv1D" in old_attn.c_attn.__class__.__name__)
            or hasattr(old_attn.c_attn, "nf")
            or (
                weight.dim() == 2
                and (weight.shape[0] < weight.shape[1] or weight.shape[1] == 3 * self.dim)
            )
        ):
            weight = weight.t().contiguous()
        split_size = weight.shape[0] // 3
        q_w = weight[:split_size, :]
        k_w = weight[split_size : 2 * split_size, :]
        v_w = weight[2 * split_size :, :]
        q_b = bias_data[:split_size] if bias_data is not None else None
        k_b = bias_data[split_size : 2 * split_size] if bias_data is not None else None
        v_b = bias_data[2 * split_size :] if bias_data is not None else None
        extracted["q"] = (q_w, q_b)
        extracted["k"] = (k_w, k_b)
        extracted["v"] = (v_w, v_b)

    if hasattr(old_attn, "c_proj") and extracted["o"][0] is None:
        w_o = old_attn.c_proj.weight.data
        if (
            ("Conv1D" in old_attn.c_proj.__class__.__name__)
            or hasattr(old_attn.c_proj, "nf")
            or (
                w_o.dim() == 2
                and (
                    w_o.shape[0] < w_o.shape[1]
                    or (w_o.shape[0] != self.dim and w_o.shape[1] == self.dim)
                )
            )
        ):
            w_o = w_o.t().contiguous()
        b_o = getattr(old_attn.c_proj, "bias", None)
        extracted["o"] = (w_o, b_o.data if b_o is not None else None)

    # 3. RoBERTa / BERT Nested Encoder handling
    if hasattr(old_attn, "self") and (
        extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
    ):
        s = old_attn.self
        if hasattr(s, "query"):
            extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
        if hasattr(s, "key"):
            extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
        if hasattr(s, "value"):
            extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

    if (
        hasattr(old_attn, "output")
        and hasattr(old_attn.output, "dense")
        and extracted["o"][0] is None
    ):
        dense = old_attn.output.dense
        extracted["o"] = (dense.weight.data, getattr(dense, "bias", None))

    # 4. Modern Decoders (Qwen3.5, LLaMA, Mistral) & DistilBERT
    named_map = {
        "q": ["q_proj", "q_lin", "query", "W_q"],
        "k": ["k_proj", "k_lin", "key", "W_k"],
        "v": ["v_proj", "v_lin", "value", "W_v"],
        "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
    }
    for proj, names in named_map.items():
        if extracted[proj][0] is None:
            for name in names:
                if hasattr(old_attn, name):
                    layer = getattr(old_attn, name)
                    base = getattr(layer, "base_layer", layer)
                    if hasattr(base, "weight"):
                        extracted[proj] = (
                            base.weight.data,
                            getattr(base, "bias", None),
                        )
                        break

    # Detect nested encoder architecture
    cls_name = old_attn.__class__.__name__
    has_out_proj = any(
        hasattr(old_attn, attr)
        for attr in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
            "out",
        )
    )
    self.is_nested_encoder = (
        cls_name in ("BertSelfAttention", "RobertaSelfAttention")
        or (cls_name.endswith("SelfAttention") and not has_out_proj)
    ) and not has_out_proj

    missing = [k for k in ["q", "k", "v"] if extracted[k][0] is None]
    if not self.is_nested_encoder and extracted["o"][0] is None:
        missing.append("o")
    if missing:
        logger.warning(
            "Attention module copy_weights warning. Proceeding with initialized weights."
        )

    def _orient(w: torch.Tensor, in_dim: int) -> torch.Tensor:
        if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim:
            return w.t().contiguous()
        return w

    with torch.no_grad():
        for proj_key, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.dim
                w_oriented = _orient(w, in_dim)
                target_w = target.weight
                out_d, in_d = target_w.shape
                src_out, src_in = w_oriented.shape
                target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                    w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                )
                if b is not None and target.bias is not None:
                    b_val = b.data if isinstance(b, nn.Parameter) else b
                    target.bias.data[: min(len(target.bias), len(b_val))].copy_(
                        b_val[: min(len(target.bias), len(b_val))]
                    )
            elif proj_key == "o" and self.is_nested_encoder:
                min_d = min(target.weight.shape[0], target.weight.shape[1])
                target.weight.data.zero_()
                target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                if target.bias is not None:
                    target.bias.data.zero_()

            # Frobenius norm verification (>= 1e-4)
            w_norm = float(torch.linalg.norm(target.weight.data.float()).item())
            if w_norm < 1e-4:
                logger.warning(
                    f"ViGAttention weight verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                )

            target.freeze_base_weights()

    # Preserve source LoRA adapters
    source_submodules = {
        "q": ["q_proj", "q_lin", "query", "W_q"],
        "k": ["k_proj", "k_lin", "key", "W_k"],
        "v": ["v_proj", "v_lin", "value", "W_v"],
        "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
    }
    for proj_key, target in [
        ("q", self.q_proj),
        ("k", self.k_proj),
        ("v", self.v_proj),
        ("o", self.out_proj),
    ]:
        for name in source_submodules[proj_key]:
            if hasattr(old_attn, name):
                sub = getattr(old_attn, name)
                if (
                    hasattr(sub, "lora_A")
                    and hasattr(sub, "lora_B")
                    and sub.lora_A is not None
                    and sub.lora_B is not None
                ):
                    src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                    src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                    if hasattr(src_lA, "data") and target.lora_A is not None:
                        if (
                            src_lA.shape[0] == target.lora_A.shape[0]
                            and src_lA.shape[1] == target.lora_A.shape[1]
                        ):
                            target.lora_A.data.copy_(src_lA.data)
                            target.lora_A.requires_grad = True
                        else:
                            print(
                                f"Warning: lora_A shape mismatch for {proj_key}; keeping zero-init."
                            )
                    if hasattr(src_lB, "data") and target.lora_B is not None:
                        if (
                            src_lB.shape[0] == target.lora_B.shape[0]
                            and src_lB.shape[1] == target.lora_B.shape[1]
                        ):
                            target.lora_B.data.copy_(src_lB.data)
                            target.lora_B.requires_grad = True
                        else:
                            print(
                                f"Warning: lora_B shape mismatch for {proj_key}; keeping zero-init."
                            )
                    if hasattr(sub, "lora_scaling"):
                        target.lora_scaling = sub.lora_scaling
                    break

    if hasattr(old_attn, "is_decoder"):
        self.is_decoder = bool(old_attn.is_decoder)

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

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

• BaseAttention:
  • extra_repr
  • forward

class XAttention (config: Any,
rank: int | None = None,
block_size: int = 32,
sparsity: float = 0.5,
lora_rank: int = 16,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "xattention", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "x_attention", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
class XAttention(BaseAttention):
    supports_causal: bool = True
    supports_bidirectional: bool = True
    supports_gqa: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True
    triton_kernel: str = "sparse_attn"

    def __init__(
        self,
        config: Any,
        rank: int | None = None,
        block_size: int = 32,
        sparsity: float = 0.5,
        lora_rank: int = 16,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        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",
                kwargs.get("num_kv_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 and recompute all head dimensions
        if "hidden_size" in kwargs:
            self.hidden_size = kwargs["hidden_size"]
        if "num_heads" in kwargs:
            self.num_heads = kwargs["num_heads"]
        if "num_kv_heads" in kwargs:
            self.num_kv_heads = kwargs["num_kv_heads"]
        if "hidden_size" in kwargs or "num_heads" in kwargs or "num_kv_heads" in kwargs:
            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.block_size: int = int(
            kwargs.get("block_size", _cfg_get(config, "xattn_block_size", block_size))
        )
        raw_keep = kwargs.get("keep_ratio", _cfg_get(config, "xattn_keep_ratio", None))
        if raw_keep is not None:
            self.keep_ratio: float = float(raw_keep)
            self.sparsity: float = 1.0 - self.keep_ratio
        else:
            raw_sparsity = float(
                kwargs.get("sparsity", _cfg_get(config, "xattn_sparsity", sparsity))
            )
            self.sparsity = raw_sparsity
            self.keep_ratio = raw_sparsity

        self.is_nested_encoder: bool = False
        self.lora_rank = kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank))

        target_dtype = kwargs.get("dtype")
        if target_dtype is None and pretrained_attn is not None:
            try:
                target_dtype = next(pretrained_attn.parameters()).dtype
            except StopIteration:
                target_dtype = torch.float32
        if target_dtype is None:
            target_dtype = torch.float32

        self.query = LoRALinear(
            self.hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.key = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.value = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
            dtype=target_dtype,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
            dtype=target_dtype,
        )

        # Trainable output scale for seamless optimizer adaptation
        self.output_scale = nn.Parameter(torch.ones(1, dtype=target_dtype))

        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(pretrained_attn, "rotary_emb", None)
        self.triton_kwargs = {"block_size": self.block_size, "eps": 1e-6}

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

        self.to(target_dtype)

    # Properties for universal aliases without polluting state_dict
    @property
    def W_q(self) -> LoRALinear:
        return self.query

    @property
    def W_k(self) -> LoRALinear:
        return self.key

    @property
    def W_v(self) -> LoRALinear:
        return self.value

    @property
    def out(self) -> LoRALinear:
        return self.out_proj

    @property
    def q_proj(self) -> LoRALinear:
        return self.query

    @property
    def k_proj(self) -> LoRALinear:
        return self.key

    @property
    def v_proj(self) -> LoRALinear:
        return self.value

    @property
    def q_lin(self) -> LoRALinear:
        return self.query

    @property
    def k_lin(self) -> LoRALinear:
        return self.key

    @property
    def v_lin(self) -> LoRALinear:
        return self.value

    @property
    def out_lin(self) -> LoRALinear:
        return self.out_proj

    @property
    def dense(self) -> LoRALinear:
        return self.out_proj

    @property
    def o_proj(self) -> LoRALinear:
        return self.out_proj

    @property
    def c_proj(self) -> LoRALinear:
        return self.out_proj

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 via instance
        adapter = UniversalAdapter()
        for proj in ["q", "k", "v", "o"]:
            try:
                w, b = adapter.get_weight_bias(src, proj)
                if w is not None:
                    extracted[proj] = (w, b)
            except Exception as e:
                logger.warning("UniversalAdapter extraction for %s failed: %s", proj, e)

        # 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_attn = getattr(src.c_attn, "base_layer", src.c_attn)
            weight = c_attn.weight.data
            bias = getattr(c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            is_conv1d = "Conv1D" in c_attn.__class__.__name__ or hasattr(c_attn, "nf")
            if is_conv1d or (weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size):
                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:
            c_proj = getattr(src.c_proj, "base_layer", src.c_proj)
            w_o = c_proj.weight.data
            is_conv1d = "Conv1D" in c_proj.__class__.__name__ or hasattr(c_proj, "nf")
            if is_conv1d or (
                w_o.dim() == 2
                and w_o.shape[0] != self.hidden_size
                and w_o.shape[1] == self.hidden_size
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(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"):
                q_mod = getattr(s.query, "base_layer", s.query)
                extracted["q"] = (q_mod.weight.data, getattr(q_mod, "bias", None))
            if hasattr(s, "key"):
                k_mod = getattr(s.key, "base_layer", s.key)
                extracted["k"] = (k_mod.weight.data, getattr(k_mod, "bias", None))
            if hasattr(s, "value"):
                v_mod = getattr(s.value, "base_layer", s.value)
                extracted["v"] = (v_mod.weight.data, getattr(v_mod, "bias", None))

        if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None:
            dense = getattr(src.output.dense, "base_layer", 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)
                        layer = getattr(layer, "base_layer", layer)
                        if hasattr(layer, "weight"):
                            extracted[proj] = (
                                layer.weight.data,
                                getattr(layer, "bias", None),
                            )
                            break

        # Detect nested encoder architecture
        cls_name = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
                "out",
            )
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ) and not has_out_proj

        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

        # Copy base weights into projection modules
        with torch.no_grad():
            for proj_key, target in [
                ("q", self.query),
                ("k", self.key),
                ("v", self.value),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape
                    target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                        w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                    )
                    if b is not None and target.bias is not None:
                        b_val = b.data if isinstance(b, nn.Parameter) else b
                        target.bias.data[: min(len(target.bias), len(b_val))].copy_(
                            b_val[: min(len(target.bias), len(b_val))]
                        )
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                # Frobenius norm verification (>= 1e-4)
                w_norm = float(torch.linalg.norm(target.weight.data.float()).item())
                if w_norm < 1e-4:
                    logger.warning(
                        f"XAttention weight copy verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                    )

                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.query),
            ("k", self.key),
            ("v", self.value),
            ("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:
                                logger.warning(
                                    "LoRA A shape mismatch for %s: src %s vs target %s; resetting",
                                    proj_key,
                                    src_lA.shape,
                                    target.lora_A.shape,
                                )
                                nn.init.kaiming_uniform_(target.lora_A, a=math.sqrt(5))
                        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:
                                logger.warning(
                                    "LoRA B shape mismatch for %s: src %s vs target %s; resetting",
                                    proj_key,
                                    src_lB.shape,
                                    target.lora_B.shape,
                                )
                                nn.init.zeros_(target.lora_B)
                        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)
        if key_value_states is not None:
            raise NotImplementedError("XAttention does not support cross-attention")
        kv_input = 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

        Q = self.query(hidden_states).view(B, N_q, H, d).transpose(1, 2).contiguous()
        K = self.key(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous()
        V = self.value(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous()

        # Apply RoPE: Priority 1 = position_embeddings, Priority 2 = rotary_emb fallback
        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:
                from neuralops.utils.compat import maybe_apply_rotary

                Q, K = maybe_apply_rotary(
                    Q,
                    K,
                    getattr(self, "adapter", None),
                    self.rotary_emb,
                    kwargs.get("position_ids"),
                )
            except Exception as e:
                logger.warning("maybe_apply_rotary failed: %s; continuing without RoPE", e)

        # GQA 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
        cfg = getattr(self, "config", None)
        is_decoder = (
            self.is_decoder
            or bool(_cfg_get(cfg, "is_decoder", False))
            or str(_cfg_get(cfg, "model_type", "")).lower()
            in ("gpt2", "qwen", "qwen2", "qwen3_5", "llama")
        )

        scores = torch.matmul(Q, K.transpose(-2, -1).contiguous()) * scale

        # True Block Sparse Routing with Padding Awareness
        if (
            self.block_size > 0
            and self.keep_ratio < 1.0
            and N_k > self.block_size
            and N_q > self.block_size
        ):
            bsz = self.block_size
            pad_q = (bsz - (N_q % bsz)) % bsz
            pad_k = (bsz - (N_k % bsz)) % bsz

            Q_pad = F.pad(Q, (0, 0, 0, pad_q)) if pad_q > 0 else Q
            K_pad = F.pad(K, (0, 0, 0, pad_k)) if pad_k > 0 else K

            total_q = Q_pad.shape[2]
            total_k = K_pad.shape[2]
            num_q_blocks = total_q // bsz
            num_k_blocks = total_k // bsz

            # Build token validity mask for Q and K to ignore padding tokens during block averaging
            q_valid = torch.zeros((B, 1, total_q, 1), device=device, dtype=Q.dtype)
            q_valid[:, :, :N_q, :] = 1.0
            k_valid = torch.zeros((B, 1, total_k, 1), device=device, dtype=K.dtype)
            k_valid[:, :, :N_k, :] = 1.0

            # If attention mask is present, factor token validity into K block averaging
            mask = attention_mask if attention_mask is not None else kwargs.get("mask")
            if mask is not None:
                valid_m = _extract_valid_token_mask(mask, B, N_k, device)
                k_valid[:, :, :N_k, 0] = k_valid[:, :, :N_k, 0] * valid_m.unsqueeze(1).to(K.dtype)

            q_valid_blocks = q_valid.view(B, 1, num_q_blocks, bsz, 1)
            Q_blocks = (Q_pad.view(B, H, num_q_blocks, bsz, d) * q_valid_blocks).sum(
                dim=3
            ) / q_valid_blocks.sum(dim=3).clamp(min=1.0)

            k_valid_blocks = k_valid.view(B, 1, num_k_blocks, bsz, 1)
            K_blocks = (K_pad.view(B, H, num_k_blocks, bsz, d) * k_valid_blocks).sum(
                dim=3
            ) / k_valid_blocks.sum(dim=3).clamp(min=1.0)

            coarse_scores = torch.matmul(Q_blocks, K_blocks.transpose(-1, -2).contiguous()) * scale

            if is_decoder:
                q_block_pos = torch.arange(
                    num_k_blocks - num_q_blocks, num_k_blocks, device=device
                ).unsqueeze(1)
                k_block_pos = torch.arange(num_k_blocks, device=device).unsqueeze(0)
                causal_block_mask = (k_block_pos <= q_block_pos).unsqueeze(0).unsqueeze(0)
                coarse_scores = coarse_scores.masked_fill(~causal_block_mask, -10000.0)

            k_blocks = max(1, min(num_k_blocks, int(math.ceil(num_k_blocks * self.keep_ratio))))
            _, topk_idx = torch.topk(coarse_scores, k=k_blocks, dim=-1)

            block_mask = torch.zeros(
                (B, H, num_q_blocks, num_k_blocks), device=device, dtype=torch.bool
            )
            block_mask.scatter_(-1, topk_idx, True)

            if is_decoder:
                block_mask = block_mask & causal_block_mask
                diag_k = min(num_q_blocks, num_k_blocks)
                diag_indices = torch.arange(diag_k, device=device)
                block_mask[:, :, diag_indices, diag_indices] = True

            # Efficiently apply block_mask to scores via broadcasting without repeat_interleave
            pad_q_total = total_q - N_q
            pad_k_total = total_k - N_k
            if pad_q_total > 0 or pad_k_total > 0:
                scores_padded = F.pad(scores, (0, pad_k_total, 0, pad_q_total), value=-10000.0)
            else:
                scores_padded = scores

            scores_padded = scores_padded.view(B, H, num_q_blocks, bsz, num_k_blocks, bsz)
            scores_padded = scores_padded.masked_fill(
                ~block_mask.unsqueeze(3).unsqueeze(5), -10000.0
            )
            scores = scores_padded.view(B, H, total_q, total_k)[:, :, :N_q, :N_k]

        # Apply causal mask
        if is_decoder:
            q_pos = torch.arange(N_k - N_q, N_k, device=device).unsqueeze(1)
            k_pos = torch.arange(N_k, device=device).unsqueeze(0)
            causal_tril = (k_pos <= q_pos).unsqueeze(0).unsqueeze(0)
            scores = scores.masked_fill(~causal_tril, -10000.0)

        # Apply 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 = _extract_valid_token_mask(mask, B, N_k, device)
                scores = scores.masked_fill((~valid_m).unsqueeze(1).unsqueeze(2), -10000.0)

        attn = torch.nan_to_num(F.softmax(scores.float(), 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:
                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)

        output = output

        if attention_mask is not None:
            valid_m = _extract_valid_token_mask(attention_mask, B, N_q, device)
            output = output * valid_m.unsqueeze(-1).to(output.dtype)

        return output, (attn if output_attentions else None)

Every 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_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.

var triton_kernel : str

The type of the None singleton.

Instance variables

prop W_k : LoRALinear

Expand source code

@property
def W_k(self) -> LoRALinear:
    return self.key

prop W_q : LoRALinear

Expand source code

@property
def W_q(self) -> LoRALinear:
    return self.query

prop W_v : LoRALinear

Expand source code

@property
def W_v(self) -> LoRALinear:
    return self.value

prop c_proj : LoRALinear

Expand source code

@property
def c_proj(self) -> LoRALinear:
    return self.out_proj

prop dense : LoRALinear

Expand source code

@property
def dense(self) -> LoRALinear:
    return self.out_proj

prop k_lin : LoRALinear

Expand source code

@property
def k_lin(self) -> LoRALinear:
    return self.key

prop k_proj : LoRALinear

Expand source code

@property
def k_proj(self) -> LoRALinear:
    return self.key

prop o_proj : LoRALinear

Expand source code

@property
def o_proj(self) -> LoRALinear:
    return self.out_proj

prop out : LoRALinear

Expand source code

@property
def out(self) -> LoRALinear:
    return self.out_proj

prop out_lin : LoRALinear

Expand source code

@property
def out_lin(self) -> LoRALinear:
    return self.out_proj

prop q_lin : LoRALinear

Expand source code

@property
def q_lin(self) -> LoRALinear:
    return self.query

prop q_proj : LoRALinear

Expand source code

@property
def q_proj(self) -> LoRALinear:
    return self.query

prop v_lin : LoRALinear

Expand source code

@property
def v_lin(self) -> LoRALinear:
    return self.value

prop v_proj : LoRALinear

Expand source code

@property
def v_proj(self) -> LoRALinear:
    return self.value

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class XAttentionScorer

Expand source code

class XAttentionScorer:
    @staticmethod
    def score_blocks(
        Q: torch.Tensor, K: torch.Tensor, block_size: int, stride: int = 2
    ) -> torch.Tensor:
        B, H, N, D = Q.shape
        n_blocks = (N + block_size - 1) // block_size

        Q_sample = Q[:, :, ::stride, :]
        K_sample = K[:, :, ::stride, :]
        N_s = Q_sample.size(2)

        scale = D**-0.5
        attn_sample = torch.matmul(Q_sample, K_sample.transpose(-2, -1).contiguous()) * scale

        b_s = max(1, block_size // stride)
        target = n_blocks * b_s
        if N_s < target:
            pad = target - N_s
            attn_sample = F.pad(attn_sample, (0, pad, 0, pad), value=0.0)
        attn_sample = attn_sample[:, :, :target, :target]

        attn_blocks = attn_sample.view(B, H, n_blocks, b_s, n_blocks, b_s)
        scores = attn_blocks.mean(dim=(3, 5))
        return scores

Static methods

def score_blocks(Q: torch.Tensor, K: torch.Tensor, block_size: int, stride: int = 2) ‑> torch.Tensor

Expand source code

@staticmethod
def score_blocks(
    Q: torch.Tensor, K: torch.Tensor, block_size: int, stride: int = 2
) -> torch.Tensor:
    B, H, N, D = Q.shape
    n_blocks = (N + block_size - 1) // block_size

    Q_sample = Q[:, :, ::stride, :]
    K_sample = K[:, :, ::stride, :]
    N_s = Q_sample.size(2)

    scale = D**-0.5
    attn_sample = torch.matmul(Q_sample, K_sample.transpose(-2, -1).contiguous()) * scale

    b_s = max(1, block_size // stride)
    target = n_blocks * b_s
    if N_s < target:
        pad = target - N_s
        attn_sample = F.pad(attn_sample, (0, pad, 0, pad), value=0.0)
    attn_sample = attn_sample[:, :, :target, :target]

    attn_blocks = attn_sample.view(B, H, n_blocks, b_s, n_blocks, b_s)
    scores = attn_blocks.mean(dim=(3, 5))
    return scores

class YourCustomAttention (config: Any | None = None,
rank: int | None = None,
pretrained_attn: torch.nn.modules.module.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("template_attention_class")
@ATTENTION_REGISTRY.register("your_custom_attention_class")
@ATTENTION_REGISTRY.register("low_rank_modulated_attention_class")
class YourCustomAttention(BaseAttention):
    r"""Template attention layer with optional low-rank pooling modulation."""

    supports_gqa: bool = True
    supports_causal: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            super().__init__(config, rank)
        except TypeError:
            super().__init__(config)

        src = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")

        # -------- hidden_size inference --------
        dim = None
        if config is not None and not isinstance(config, int):
            dim = _cfg_get(
                config,
                "hidden_size",
                _cfg_get(config, "dim", _cfg_get(config, "d_model", None)),
            )
        if dim is None and src is not None:
            dim = getattr(src, "hidden_size", getattr(src, "embed_dim", None))
            if dim is None and hasattr(src, "q_proj") and hasattr(src.q_proj, "in_features"):
                dim = src.q_proj.in_features
            elif dim is None and hasattr(src, "q_lin") and hasattr(src.q_lin, "in_features"):
                dim = src.q_lin.in_features
            elif dim is None and hasattr(src, "query") and hasattr(src.query, "in_features"):
                dim = src.query.in_features
            elif dim is None and hasattr(src, "c_attn") and hasattr(src.c_attn, "weight"):
                w = src.c_attn.weight
                dim = w.shape[0] if w.shape[0] < w.shape[1] else w.shape[1]
        if dim is None:
            dim = kwargs.get("hidden_size", 768)
        self.hidden_size: int = int(dim)

        # -------- num_heads inference (TA-13) --------
        num_heads = None
        if config is not None and not isinstance(config, int):
            num_heads = _cfg_get(
                config,
                "num_attention_heads",
                _cfg_get(config, "n_heads", _cfg_get(config, "num_heads", None)),
            )
        if num_heads is None and src is not None:
            num_heads = getattr(
                src,
                "num_heads",
                getattr(src, "num_attention_heads", getattr(src, "n_head", None)),
            )
        if num_heads is None:
            # Fallback: assume head_dim ≈ 64
            num_heads = max(1, self.hidden_size // 64)
        self.num_heads: int = max(1, int(num_heads))

        # -------- head_dim --------
        head_dim = None
        if config is not None and not isinstance(config, int):
            head_dim = _cfg_get(config, "head_dim", None)
        if head_dim is None:
            head_dim = kwargs.get("head_dim")
        if head_dim is None:
            head_dim = self.hidden_size // self.num_heads
        self.head_dim: int = int(head_dim)

        # -------- num_kv_heads (GQA) --------
        cfg_kv = None
        if config is not None and not isinstance(config, int):
            cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None))
        if cfg_kv is None and src is not None:
            cfg_kv = getattr(src, "num_key_value_heads", getattr(src, "num_kv_heads", None))
        self.num_kv_heads: int = int(cfg_kv) if cfg_kv is not None else self.num_heads
        self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads))
        self.num_key_value_groups: int = max(1, self.num_heads // self.num_kv_heads)

        # TA-04: all_head_size explicitly defined
        self.q_head_size: int = self.num_heads * self.head_dim
        self.kv_head_size: int = self.num_kv_heads * self.head_dim
        self.all_head_size: int = self.q_head_size

        self.lora_rank: int = int(
            kwargs.get(
                "lora_rank",
                (_cfg_get(config, "lora_rank", lora_rank) if config is not None else lora_rank),
            )
        )
        self.lora_alpha: float = float(
            kwargs.get(
                "lora_alpha",
                (_cfg_get(config, "lora_alpha", lora_alpha) if config is not None else lora_alpha),
            )
        )

        # -------- decoder detection --------
        if is_decoder is not None:
            self.is_decoder: bool = bool(is_decoder)
        elif config is not None and not isinstance(config, int):
            mt = str(_cfg_get(config, "model_type", "") or "").lower()
            self.is_decoder = (
                bool(_cfg_get(config, "is_decoder", False))
                or bool(kwargs.get("is_causal", False))
                or any(
                    mt.startswith(p)
                    for p in (
                        "gpt",
                        "qwen",
                        "llama",
                        "mistral",
                        "gemma",
                        "phi",
                        "opt",
                        "bloom",
                        "falcon",
                        "olmo",
                        "mpt",
                        "stablelm",
                    )
                )
            )
        else:
            self.is_decoder = bool(kwargs.get("is_causal", False))
        self.is_causal = self.is_decoder
        self.is_nested_encoder: bool = False

        # -------- projections --------
        self.query = LoRALinear(
            self.hidden_size,
            self.q_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.key = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.value = 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.q_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        # Aliases
        self.q_proj = self.query
        self.k_proj = self.key
        self.v_proj = self.value

        # TA-06: rank_proj respects device/dtype
        if rank is not None and rank > 0:
            self.rank_proj: nn.Linear | None = nn.Linear(
                self.hidden_size,
                rank,
                bias=False,
            )
            nn.init.zeros_(self.rank_proj.weight)
        else:
            self.rank_proj = None

        dropout_p = 0.1
        if config is not None and not isinstance(config, int):
            dropout_p = _cfg_get(
                config, "dropout", _cfg_get(config, "attention_probs_dropout_prob", 0.1)
            )
        self.dropout = nn.Dropout(dropout_p)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    # ------------------------------------------------------------------
    @property
    def out_lin(self):
        return self.out_proj

    @property
    def o_proj(self):
        return self.out_proj

    @property
    def dense(self):
        return self.out_proj

    @property
    def c_proj(self):
        return self.out_proj

    @property
    def q_lin(self):
        return self.query

    @property
    def k_lin(self):
        return self.key

    @property
    def v_lin(self):
        return self.value

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2, Qwen/LLaMA."""
        with torch.no_grad():
            w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

            src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
            if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
                hasattr(old_attn, "query")
                and not hasattr(old_attn, "dense")
                and not hasattr(old_attn, "out_proj")
            ):
                self.is_nested_encoder = True

            if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
                self.is_nested_encoder = True
                src = old_attn.self
                if hasattr(old_attn.output, "dense"):
                    w_o = old_attn.output.dense.weight.detach()
                    b_o = (
                        old_attn.output.dense.bias.detach()
                        if old_attn.output.dense.bias is not None
                        else None
                    )
            else:
                src = old_attn

            # Pattern 1: GPT-2 Conv1D fused c_attn
            if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
                c_attn_w = src.c_attn.weight.detach()
                c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
                if c_attn_w.shape[0] == self.hidden_size and c_attn_w.shape[1] != self.hidden_size:
                    c_attn_w = c_attn_w.t().contiguous()
                q_dim, k_dim, v_dim = (
                    self.q_head_size,
                    self.kv_head_size,
                    self.kv_head_size,
                )
                w_q = c_attn_w[:q_dim, :]
                w_k = c_attn_w[q_dim : q_dim + k_dim, :]
                w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
                if c_attn_b is not None:
                    b_q = c_attn_b[:q_dim]
                    b_k = c_attn_b[q_dim : q_dim + k_dim]
                    b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]
                out_w = src.c_proj.weight.detach()
                if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                    out_w = out_w.t().contiguous()
                w_o = out_w
                b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

            # Pattern 2: RoBERTa / BERT
            elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
                w_q = src.query.weight.detach()
                b_q = src.query.bias.detach() if src.query.bias is not None else None
                w_k = src.key.weight.detach()
                b_k = src.key.bias.detach() if src.key.bias is not None else None
                w_v = src.value.weight.detach()
                b_v = src.value.bias.detach() if src.value.bias is not None else None
                if w_o is None:
                    for attr in ("dense", "out_proj", "o_proj"):
                        if hasattr(src, attr):
                            proj = getattr(src, attr)
                            w_o = proj.weight.detach()
                            b_o = proj.bias.detach() if proj.bias is not None else None
                            break

            # Pattern 3: Qwen / LLaMA
            elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
                w_q = src.q_proj.weight.detach()
                b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
                w_k = src.k_proj.weight.detach()
                b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
                w_v = src.v_proj.weight.detach()
                b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
                if hasattr(src, "o_proj"):
                    w_o = src.o_proj.weight.detach()
                    b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

            # Pattern 4: DistilBERT
            elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
                w_q = src.q_lin.weight.detach()
                b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
                w_k = src.k_lin.weight.detach()
                b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
                w_v = src.v_lin.weight.detach()
                b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
                if hasattr(src, "out_lin"):
                    w_o = src.out_lin.weight.detach()
                    b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

            # Fallback: UniversalAdapter
            if w_q is None or w_k is None or w_v is None:
                try:
                    adapter = UniversalAdapter()
                    if w_q is None:
                        w_q, b_q = adapter.get_weight_bias(src, "q")
                    if w_k is None:
                        w_k, b_k = adapter.get_weight_bias(src, "k")
                    if w_v is None:
                        w_v, b_v = adapter.get_weight_bias(src, "v")
                    if w_o is None:
                        w_o, b_o = adapter.get_weight_bias(src, "o")
                except Exception as e:
                    logger.debug("UniversalAdapter fallback failed: %s", e)

            missing = []
            if w_q is None:
                missing.append("Q")
            if w_k is None:
                missing.append("K")
            if w_v is None:
                missing.append("V")
            if w_o is None and not self.is_nested_encoder:
                missing.append("O")
            if missing:
                raise RuntimeError(
                    f"YourCustomAttention.copy_weights: missing {missing} from {type(old_attn).__name__}"
                )

            def _copy_layer(target_layer: LoRALinear, w, b) -> None:
                if w is not None:
                    if w.shape == target_layer.weight.shape:
                        target_layer.weight.data.copy_(w.to(target_layer.weight.dtype))
                    elif w.t().contiguous().shape == target_layer.weight.shape:
                        target_layer.weight.data.copy_(
                            w.t().contiguous().to(target_layer.weight.dtype)
                        )
                    else:
                        min_0 = min(target_layer.weight.shape[0], w.shape[0])
                        min_1 = min(target_layer.weight.shape[1], w.shape[1])
                        target_layer.weight.data.zero_()
                        target_layer.weight.data[:min_0, :min_1].copy_(
                            w[:min_0, :min_1].to(target_layer.weight.dtype)
                        )

                if b is not None and target_layer.bias is not None:
                    if b.shape == target_layer.bias.shape:
                        target_layer.bias.data.copy_(b.to(target_layer.bias.dtype))
                    else:
                        min_b = min(target_layer.bias.shape[0], b.shape[0])
                        target_layer.bias.data.zero_()
                        target_layer.bias.data[:min_b].copy_(b[:min_b].to(target_layer.bias.dtype))
                elif target_layer.bias is not None:
                    target_layer.bias.data.zero_()

            _copy_layer(self.query, w_q, b_q)
            _copy_layer(self.key, w_k, b_k)
            _copy_layer(self.value, w_v, b_v)
            if w_o is not None:
                _copy_layer(self.out_proj, w_o, b_o)

            # Frobenius validation
            for name, layer in (
                ("q", self.query),
                ("k", self.key),
                ("v", self.value),
                ("o", self.out_proj),
            ):
                norm = float(torch.linalg.norm(layer.weight.data.float()).item())
                if norm < 1e-4:
                    logger.warning(
                        "YourCustomAttention: '%s' weight norm near zero (%.2e)",
                        name,
                        norm,
                    )

    # ------------------------------------------------------------------
    # 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,
    ):
        if encoder_hidden_states is not None:
            raise NotImplementedError("YourCustomAttention: cross-attention not supported.")

        B, L, _ = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        q = (
            self.query(hidden_states)
            .view(B, L, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            self.key(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            self.value(hidden_states)
            .view(B, L, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # TA-02: RoPE with past_len offset
        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]
            cos = cos[..., past_len : past_len + L, :]
            sin = sin[..., past_len : past_len + L, :]
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # GQA expansion
        if self.num_kv_heads != self.num_heads:
            k = repeat_kv(k, self.num_key_value_groups)
            v = repeat_kv(v, self.num_key_value_groups)

        scale = 1.0 / math.sqrt(self.head_dim)

        # TA-05: unify mask
        eff_mask = encoder_attention_mask if encoder_attention_mask is not None else attention_mask
        if eff_mask is None:
            eff_mask = kwargs.get("attn_mask")

        valid_mask_2d = None
        if eff_mask is not None:
            if eff_mask.dim() == 2:
                if eff_mask.dtype == torch.bool:
                    valid_mask_2d = eff_mask
                elif eff_mask.dtype.is_floating_point:
                    valid_mask_2d = eff_mask > -1e4
                else:
                    valid_mask_2d = eff_mask > 0
            elif eff_mask.dim() == 4:
                if eff_mask.dtype == torch.bool:
                    v_mask = eff_mask
                elif eff_mask.dtype.is_floating_point:
                    v_mask = eff_mask > -1e4
                else:
                    v_mask = eff_mask > 0
                valid_mask_2d = v_mask.any(dim=-2).squeeze(1)

        # SDPA fast path (TA-03: mask-aware)
        is_causal = self.is_decoder and L > 1
        use_sdpa = self.rank_proj is None and head_mask is None and not output_attentions

        if use_sdpa:
            try:
                q_sdpa = q.half() if q.dtype == torch.float32 else q
                k_sdpa = k.half() if k.dtype == torch.float32 else k
                v_sdpa = v.half() if v.dtype == torch.float32 else v

                attn_bias = None
                if eff_mask is not None:
                    if eff_mask.dim() == 2:
                        attn_bias = (
                            (eff_mask.unsqueeze(1).unsqueeze(2).float() - 1.0) * 10000.0
                        ).to(q_sdpa.dtype)
                    elif eff_mask.dim() == 4:
                        attn_bias = eff_mask.to(q_sdpa.dtype)

                out_sdpa = F.scaled_dot_product_attention(
                    q_sdpa,
                    k_sdpa,
                    v_sdpa,
                    attn_mask=attn_bias,
                    is_causal=is_causal and attn_bias is None,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                out = out_sdpa.to(dtype).transpose(1, 2).contiguous().view(B, L, self.q_head_size)
                out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)

                if self.is_nested_encoder:
                    return out, None
                result = self.out_proj(out)
                result = torch.nan_to_num(result, nan=0.0, posinf=1e4, neginf=-1e4)
                return result, None
            except Exception as e:
                # TA-08: log instead of silent pass
                logger.debug("SDPA path failed (%s); falling back to explicit.", e)

        # Explicit path
        scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale

        # Low-rank modulation
        if self.rank_proj is not None:
            H_pool = hidden_states.mean(dim=1)
            H_rank = self.rank_proj(H_pool)
            modulation = 1.0 + torch.tanh(H_rank).mean(dim=-1, keepdim=True).view(B, 1, 1, 1)
            scores = scores * modulation

        if is_causal:
            causal_mask = (
                torch.triu(
                    torch.full((L, L), float("-inf"), device=device, dtype=torch.float32),
                    diagonal=1,
                )
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores = scores + causal_mask

        if valid_mask_2d is not None:
            scores = scores.masked_fill(~valid_mask_2d[:, None, None, :], -1e4)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(dtype)
        attn = torch.nan_to_num(attn, nan=0.0)
        attn = self.dropout(attn)

        # TA-05: head_mask unified
        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)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)
            attn = attn * head_mask.to(attn.dtype)

        context = torch.matmul(attn, v)
        context = context.transpose(1, 2).contiguous().view(B, L, self.q_head_size)
        context = torch.nan_to_num(context, nan=0.0, posinf=1e4, neginf=-1e4)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out_proj(context)

        # TA-14: no silent pad/truncate
        if output.shape[-1] != self.hidden_size:
            raise RuntimeError(
                f"YourCustomAttention: output dim {output.shape[-1]} != hidden_size {self.hidden_size}"
            )

        return output, (attn if output_attentions else None)

Template attention layer with optional low-rank pooling modulation.

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.

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.key

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.query

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.value

Methods

def copy_weights(self,
old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2, Qwen/LLaMA."""
    with torch.no_grad():
        w_q = b_q = w_k = b_k = w_v = b_v = w_o = b_o = None

        src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__
        if src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (
            hasattr(old_attn, "query")
            and not hasattr(old_attn, "dense")
            and not hasattr(old_attn, "out_proj")
        ):
            self.is_nested_encoder = True

        if hasattr(old_attn, "self") and hasattr(old_attn, "output"):
            self.is_nested_encoder = True
            src = old_attn.self
            if hasattr(old_attn.output, "dense"):
                w_o = old_attn.output.dense.weight.detach()
                b_o = (
                    old_attn.output.dense.bias.detach()
                    if old_attn.output.dense.bias is not None
                    else None
                )
        else:
            src = old_attn

        # Pattern 1: GPT-2 Conv1D fused c_attn
        if hasattr(src, "c_attn") and hasattr(src, "c_proj"):
            c_attn_w = src.c_attn.weight.detach()
            c_attn_b = src.c_attn.bias.detach() if src.c_attn.bias is not None else None
            if c_attn_w.shape[0] == self.hidden_size and c_attn_w.shape[1] != self.hidden_size:
                c_attn_w = c_attn_w.t().contiguous()
            q_dim, k_dim, v_dim = (
                self.q_head_size,
                self.kv_head_size,
                self.kv_head_size,
            )
            w_q = c_attn_w[:q_dim, :]
            w_k = c_attn_w[q_dim : q_dim + k_dim, :]
            w_v = c_attn_w[q_dim + k_dim : q_dim + k_dim + v_dim, :]
            if c_attn_b is not None:
                b_q = c_attn_b[:q_dim]
                b_k = c_attn_b[q_dim : q_dim + k_dim]
                b_v = c_attn_b[q_dim + k_dim : q_dim + k_dim + v_dim]
            out_w = src.c_proj.weight.detach()
            if out_w.shape[0] == self.hidden_size and out_w.shape[1] != self.hidden_size:
                out_w = out_w.t().contiguous()
            w_o = out_w
            b_o = src.c_proj.bias.detach() if src.c_proj.bias is not None else None

        # Pattern 2: RoBERTa / BERT
        elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"):
            w_q = src.query.weight.detach()
            b_q = src.query.bias.detach() if src.query.bias is not None else None
            w_k = src.key.weight.detach()
            b_k = src.key.bias.detach() if src.key.bias is not None else None
            w_v = src.value.weight.detach()
            b_v = src.value.bias.detach() if src.value.bias is not None else None
            if w_o is None:
                for attr in ("dense", "out_proj", "o_proj"):
                    if hasattr(src, attr):
                        proj = getattr(src, attr)
                        w_o = proj.weight.detach()
                        b_o = proj.bias.detach() if proj.bias is not None else None
                        break

        # Pattern 3: Qwen / LLaMA
        elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"):
            w_q = src.q_proj.weight.detach()
            b_q = src.q_proj.bias.detach() if src.q_proj.bias is not None else None
            w_k = src.k_proj.weight.detach()
            b_k = src.k_proj.bias.detach() if src.k_proj.bias is not None else None
            w_v = src.v_proj.weight.detach()
            b_v = src.v_proj.bias.detach() if src.v_proj.bias is not None else None
            if hasattr(src, "o_proj"):
                w_o = src.o_proj.weight.detach()
                b_o = src.o_proj.bias.detach() if src.o_proj.bias is not None else None

        # Pattern 4: DistilBERT
        elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"):
            w_q = src.q_lin.weight.detach()
            b_q = src.q_lin.bias.detach() if src.q_lin.bias is not None else None
            w_k = src.k_lin.weight.detach()
            b_k = src.k_lin.bias.detach() if src.k_lin.bias is not None else None
            w_v = src.v_lin.weight.detach()
            b_v = src.v_lin.bias.detach() if src.v_lin.bias is not None else None
            if hasattr(src, "out_lin"):
                w_o = src.out_lin.weight.detach()
                b_o = src.out_lin.bias.detach() if src.out_lin.bias is not None else None

        # Fallback: UniversalAdapter
        if w_q is None or w_k is None or w_v is None:
            try:
                adapter = UniversalAdapter()
                if w_q is None:
                    w_q, b_q = adapter.get_weight_bias(src, "q")
                if w_k is None:
                    w_k, b_k = adapter.get_weight_bias(src, "k")
                if w_v is None:
                    w_v, b_v = adapter.get_weight_bias(src, "v")
                if w_o is None:
                    w_o, b_o = adapter.get_weight_bias(src, "o")
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        missing = []
        if w_q is None:
            missing.append("Q")
        if w_k is None:
            missing.append("K")
        if w_v is None:
            missing.append("V")
        if w_o is None and not self.is_nested_encoder:
            missing.append("O")
        if missing:
            raise RuntimeError(
                f"YourCustomAttention.copy_weights: missing {missing} from {type(old_attn).__name__}"
            )

        def _copy_layer(target_layer: LoRALinear, w, b) -> None:
            if w is not None:
                if w.shape == target_layer.weight.shape:
                    target_layer.weight.data.copy_(w.to(target_layer.weight.dtype))
                elif w.t().contiguous().shape == target_layer.weight.shape:
                    target_layer.weight.data.copy_(
                        w.t().contiguous().to(target_layer.weight.dtype)
                    )
                else:
                    min_0 = min(target_layer.weight.shape[0], w.shape[0])
                    min_1 = min(target_layer.weight.shape[1], w.shape[1])
                    target_layer.weight.data.zero_()
                    target_layer.weight.data[:min_0, :min_1].copy_(
                        w[:min_0, :min_1].to(target_layer.weight.dtype)
                    )

            if b is not None and target_layer.bias is not None:
                if b.shape == target_layer.bias.shape:
                    target_layer.bias.data.copy_(b.to(target_layer.bias.dtype))
                else:
                    min_b = min(target_layer.bias.shape[0], b.shape[0])
                    target_layer.bias.data.zero_()
                    target_layer.bias.data[:min_b].copy_(b[:min_b].to(target_layer.bias.dtype))
            elif target_layer.bias is not None:
                target_layer.bias.data.zero_()

        _copy_layer(self.query, w_q, b_q)
        _copy_layer(self.key, w_k, b_k)
        _copy_layer(self.value, w_v, b_v)
        if w_o is not None:
            _copy_layer(self.out_proj, w_o, b_o)

        # Frobenius validation
        for name, layer in (
            ("q", self.query),
            ("k", self.key),
            ("v", self.value),
            ("o", self.out_proj),
        ):
            norm = float(torch.linalg.norm(layer.weight.data.float()).item())
            if norm < 1e-4:
                logger.warning(
                    "YourCustomAttention: '%s' weight norm near zero (%.2e)",
                    name,
                    norm,
                )

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2, Qwen/LLaMA.

Inherited members

• BaseAttention:
  • extra_repr
  • forward

class ZeroSAttention (config: Any,
rank: int | None = None,
lora_rank: int = 16,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "zeros", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "zeros_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class ZeroSAttention(BaseAttention):
    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,
        rank: int | None = None,
        lora_rank: int = 16,
        pretrained_attn: nn.Module | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        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))

        self.is_nested_encoder: bool = False
        self.lora_rank = _cfg_get(config, "lora_rank", lora_rank)

        self.W_q = LoRALinear(
            self.hidden_size,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.W_k = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.W_v = LoRALinear(
            self.hidden_size,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )
        self.out = LoRALinear(
            self.all_head_size,
            self.hidden_size,
            bias=True,
            lora_rank=self.lora_rank,
            freeze_base=True,
        )

        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(pretrained_attn, "rotary_emb", None)

        if pretrained_attn is not None:
            self.copy_weights(pretrained_attn)

    def copy_weights(self, src: Any) -> None:
        from neuralops.core.adapters import UniversalAdapter

        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 as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # 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
        ):
            weight = src.c_attn.weight.data
            bias = getattr(src.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if weight.dim() == 2 and weight.shape[1] == 3 * self.hidden_size:
                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:
            w_o = src.c_proj.weight.data
            if (
                w_o.dim() == 2
                and w_o.shape[0] != self.hidden_size
                and w_o.shape[1] == self.hidden_size
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(src.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(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"],
        }
        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 = src.__class__.__name__
        has_out_proj = any(
            hasattr(src, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
                "out",
            )
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ) and not has_out_proj

        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

        # Copy base weights into projection modules
        with torch.no_grad():
            for proj_key, target in [
                ("q", self.W_q),
                ("k", self.W_k),
                ("v", self.W_v),
                ("o", self.out),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.hidden_size
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape
                    target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                        w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                    )
                    if b is not None and target.bias is not None:
                        b_val = b.data if isinstance(b, nn.Parameter) else b
                        target.bias.data[: min(len(target.bias), len(b_val))].copy_(
                            b_val[: min(len(target.bias), len(b_val))]
                        )
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                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.W_q),
            ("k", self.W_k),
            ("v", self.W_v),
            ("o", self.out),
        ]:
            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, L, D = hidden_states.shape
        N_k = kv_input.shape[1]
        H, Dh = self.num_heads, self.head_dim
        device = hidden_states.device

        Q = self.W_q(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous()
        K = self.W_k(kv_input).view(B, N_k, self.num_kv_heads, Dh).transpose(1, 2).contiguous()
        V = self.W_v(kv_input).view(B, N_k, self.num_kv_heads, Dh).transpose(1, 2).contiguous()

        # Apply RoPE: Priority 1 = position_embeddings, Priority 2 = rotary_emb fallback
        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:
                    print(f"copy_weights failed: {e}")
                    logger.warning(f"copy_weights failed: {e}")

        # GQA repeat
        if self.kv_repeat > 1:
            K = repeat_kv(K, self.kv_repeat)
            V = repeat_kv(V, self.kv_repeat)

        is_decoder = (
            self.is_decoder
            or getattr(self.config, "is_decoder", False)
            or getattr(self.config, "model_type", "")
            in ("gpt2", "qwen", "qwen2", "qwen3_5", "llama")
        )

        eps = 1e-6
        attn_map = None
        if is_decoder:
            valid_m = get_token_valid_mask(attention_mask, B, N_k, device)
            if valid_m is not None:
                mask_f = valid_m.unsqueeze(1).unsqueeze(-1).to(K.dtype)
                K = K * mask_f
                V = V * mask_f
                k_cum_sum = torch.cumsum(K, dim=2)
                count = torch.cumsum(mask_f, dim=2).clamp(min=1.0)
                k_running_mean = k_cum_sum / count
            else:
                step = torch.arange(1, N_k + 1, device=device, dtype=K.dtype).view(1, 1, N_k, 1)
                k_cum_sum = torch.cumsum(K, dim=2)
                k_running_mean = k_cum_sum / step

            k_centered = K - k_running_mean

            q_hat = F.relu(Q) + eps  # [B, H, L, d]
            k_hat = F.relu(k_centered) + eps  # [B, H, N_k, d]

            # Causal linear attention via cumulative states
            kv = torch.einsum("b h j d, b h j e -> b h j d e", k_hat, V)  # [B, H, N_k, d, d]
            kv_cum = torch.cumsum(kv, dim=2)  # [B, H, N_k, d, d]
            k_cum = torch.cumsum(k_hat, dim=2)  # [B, H, N_k, d]

            sim = None
            if N_k == L:
                num = torch.einsum("b h l d, b h l d e -> b h l e", q_hat, kv_cum)  # [B, H, L, d]
                den = (q_hat * k_cum).sum(dim=-1, keepdim=True)  # [B, H, L, 1]
            elif N_k >= L:
                kv_cum_slice = kv_cum[:, :, -L:, :, :]
                k_cum_slice = k_cum[:, :, -L:, :]
                num = torch.einsum(
                    "b h l d, b h l d e -> b h l e", q_hat, kv_cum_slice
                )  # [B, H, L, d]
                den = (q_hat * k_cum_slice).sum(dim=-1, keepdim=True)  # [B, H, L, 1]
            else:
                sim = torch.matmul(q_hat, k_hat.transpose(-1, -2).contiguous())  # [B, H, L, N_k]
                causal_mask = torch.tril(torch.ones(L, N_k, device=device, dtype=torch.bool))
                sim = sim.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), 0.0)
                num = torch.matmul(sim, V)  # [B, H, L, d]
                den = sim.sum(dim=-1, keepdim=True)  # [B, H, L, 1]

            context = num / den.clamp(min=eps)  # [B, H, L, d]
            context = self.dropout(context)

            if output_attentions:
                if sim is not None:
                    attn_map = sim / den.clamp(min=eps)
                else:
                    attn_map = None
        else:
            # Bidirectional Mean-Centering
            valid_m = get_token_valid_mask(attention_mask, B, N_k, device)
            if valid_m is not None:
                mask_f = valid_m.unsqueeze(1).unsqueeze(-1).to(K.dtype)
                K = K * mask_f
                k_mean = K.sum(dim=2, keepdim=True) / mask_f.sum(dim=2, keepdim=True).clamp(min=1.0)
            else:
                k_mean = K.mean(dim=2, keepdim=True)

            k_centered = K - k_mean
            q_hat = F.relu(Q) + eps
            k_hat = F.relu(k_centered) + eps

            S_kv = torch.matmul(k_hat.transpose(-2, -1).contiguous(), V)  # [B, H, d, d]
            num = torch.matmul(q_hat, S_kv)  # [B, H, L, d]
            den = (
                torch.matmul(q_hat, k_hat.sum(dim=2, keepdim=True).transpose(-2, -1).contiguous())
                + eps
            )
            context = num / den.clamp(min=eps)
            context = self.dropout(context)

            if output_attentions:
                attn_map = torch.matmul(q_hat, k_hat.transpose(-1, -2).contiguous()) / den.clamp(
                    min=eps
                )

        # 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:
                if (
                    head_mask.shape[-1] == head_mask.shape[-2]
                    and head_mask.shape[-1] != context.shape[-1]
                ):
                    head_mask = head_mask.mean(dim=-1, keepdim=True)
            context = context * head_mask.to(context.dtype)

        context = context.transpose(1, 2).contiguous().view(B, L, self.all_head_size)

        if self.is_nested_encoder:
            output = context
        else:
            output = self.out(context)

        return output, attn_map

Every 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_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.

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward

class SDPBackend (...)

An enum-like class that contains the different backends for scaled dot product attention.

… warning:: This class is in beta and subject to change.

This backend class is designed to be used with the sdpa_kernel context manager.See :func: torch.nn.attention.sdpa_kernel for more details.

Members -----=

ERROR

MATH

FLASH_ATTENTION

EFFICIENT_ATTENTION

CUDNN_ATTENTION

OVERRIDEABLE

init(self: torch._C._SDPBackend, value: typing.SupportsInt) -> None

Ancestors

• pybind11_builtins.pybind11_object

Class variables

var CUDNN_ATTENTION

The type of the None singleton.

var EFFICIENT_ATTENTION

The type of the None singleton.

var ERROR

The type of the None singleton.

var FLASH_ATTENTION

The type of the None singleton.

var MATH

The type of the None singleton.

var OVERRIDEABLE

The type of the None singleton.

Instance variables

prop name : str

prop value